Cast the lm_head weight to the compute dtype in the chunked projections

Under FSDP2 mixed precision the patched forward reads the fp32 sharded master
weight directly while the hidden states come out bf16, so the tensor-core
projection failed with a dtype mismatch (tests/distributed test_sft_peft[fsdp2]).
Cast the weight to the hidden-states dtype, which is what lm_head's own forward
computes in. Same fix in all three copies: SFT, distillation, async distillation.

Author: qgallouedec

trl/experimental/async_distillation/async_distillation_trainer.py

@@ -225,8 +225,9 @@ def _jsd_loss_chunk(
         are `torch.no_grad()` sums for this chunk only — callers accumulate them across chunks and reduce across ranks,
         exactly as the non-chunked path already did.
     """
-    # Project in the model dtype and upcast only afterwards, as the other chunked projections do.
-    logits = (hidden_chunk @ lm_head_weight.t()).float()
+    # Project in the compute dtype and upcast only afterwards, as the other chunked projections do, casting
+    # the weight to the hidden-states dtype in case it is an fp32 master weight (FSDP2 mixed precision).
+    logits = (hidden_chunk @ lm_head_weight.to(hidden_chunk.dtype).t()).float()
     if lm_head_bias is not None:
         logits = logits + lm_head_bias.float()
     if logit_scale != 1.0:

trl/trainer/distillation_trainer.py

@@ -112,9 +112,11 @@ def _chunk(h_s, w_s, b_s, s_scale, s_softcap, h_t, w_t, b_t, t_scale, t_softcap,
     # the backward, never `(chunk, V)`. ZeRO-3 shards the `lm_head`, so gather it tightly around each projection.
     # `logit_scale` (Cohere) / `final_logit_softcapping` (Gemma) are applied per model to match its full forward.
     with maybe_gather_lm_head_ctx(w_s, b_s):
-        # Project in the model dtype and upcast only afterwards, as `"nll"` and `transformers`'
-        # `ForCausalLMLoss` do.
-        student_logits = (h_s @ w_s.t()).float()
+        # Project in the compute dtype and upcast only afterwards, as `"nll"` and `transformers`'
+        # `ForCausalLMLoss` do. The weight can be an fp32 master weight while the hidden states are
+        # bf16 (FSDP2 mixed precision reads the sharded weight directly, bypassing the cast its
+        # forward hooks apply), so cast it to the dtype `lm_head`'s own forward would compute in.
+        student_logits = (h_s @ w_s.to(h_s.dtype).t()).float()
         if b_s is not None:
             student_logits = student_logits + b_s.float()
     if s_scale != 1.0:
@@ -125,7 +127,7 @@ def _chunk(h_s, w_s, b_s, s_scale, s_softcap, h_t, w_t, b_t, t_scale, t_softcap,
     # and the teacher accumulates no gradients (the teacher params are not frozen by `prepare_model`). Everything
     # downstream inherits this since `teacher_logits` is already detached.
     with maybe_gather_lm_head_ctx(w_t, b_t), torch.no_grad():
-        teacher_logits = (h_t @ w_t.t()).float()
+        teacher_logits = (h_t @ w_t.to(h_t.dtype).t()).float()
         if b_t is not None:
             teacher_logits = teacher_logits + b_t.float()
     if t_scale != 1.0:

trl/trainer/sft_trainer.py

@@ -99,9 +99,11 @@ class _ChunkedCELMHeadOutput(CausalLMOutputWithPast):
 
 def _chunk(h, w, b, lbl, logit_scale, final_logit_softcapping):
     with maybe_gather_lm_head_ctx(w, b):
-        # Project in the model dtype and upcast only for the softmax, like `"nll"` and
-        # `transformers`' own `ForCausalLMLoss` do.
-        logits = (h @ w.t()).float()
+        # Project in the compute dtype and upcast only for the softmax, like `"nll"` and
+        # `transformers`' own `ForCausalLMLoss` do. `w` can be an fp32 master weight while `h` is
+        # bf16 (FSDP2 mixed precision reads the sharded weight directly, bypassing the cast its
+        # forward hooks apply), so cast it to the dtype `lm_head`'s own forward would compute in.
+        logits = (h @ w.to(h.dtype).t()).float()
         if b is not None:
             logits = logits + b.float()
     if logit_scale != 1.0: