Move logits.float() call

[ringohoffman]

Summary

The analogous logits.float() calls were moved in the Hugging Face modeling source code to be inside the if labels is not None block to avoid upcasting logits unless they are being used in a loss calculation; this avoids a memory spike during inference if the model is in lower precision.

• https://github.com/huggingface/transformers/blob/37ea04013b34b39c01b51aeaacd8d56f2c62a7eb/src/transformers/models/llama/modeling_llama.py#L1211-L1212
• https://github.com/huggingface/transformers/blob/37ea04013b34b39c01b51aeaacd8d56f2c62a7eb/src/transformers/models/mixtral/modeling_mixtral.py#L1329-L1330
• https://github.com/huggingface/transformers/blob/37ea04013b34b39c01b51aeaacd8d56f2c62a7eb/src/transformers/models/phi3/modeling_phi3.py#L1303-L1304
• https://github.com/huggingface/transformers/blob/37ea04013b34b39c01b51aeaacd8d56f2c62a7eb/src/transformers/models/qwen2/modeling_qwen2.py#L1206-L1207

Some of your models already have this change:

Liger-Kernel/src/liger_kernel/transformers/model/mistral.py

Lines 114 to 116 in ff6650b

	if labels is not None:
	# Upcast to float if we need to compute the loss to avoid potential precision issues
	logits = logits.float()

Liger-Kernel/src/liger_kernel/transformers/model/gemma.py

Lines 114 to 116 in ff6650b

	if labels is not None:
	# Upcast to float if we need to compute the loss to avoid potential precision issues
	logits = logits.float()

See also:

• Improving memory efficiency further 🚀 huggingface/transformers#30860

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[ringohoffman]

@shimizust @lancerts @ByronHsu

How does this look?

[ByronHsu]

Sorry for the delay. This looks great! Thank you for making liger consistent with the latest hf implementation