Do the chunked-CE lm_head projection on tensor cores instead of in fp32

[qgallouedec]

_chunk in sft_trainer.py, the inner loop of the default loss_type="chunked_nll" does

trl/trl/trainer/sft_trainer.py

Line 101 in 6d484ba

logits = h.float() @ w.float().t()

Both operands are already the model dtype (bf16), so upcasting them buys no information; it just moves the GEMM off the tensor cores onto the fp32 SIMT path and materialises an fp32 copy of the entire lm_head weight. For a 248,320-token vocabulary that copy is 2.03 GB, rebuilt for every chunk and again on every gradient-checkpoint recompute.

logits = (h @ w.t()).float()

is what loss_type="nll" already gets (the model applies lm_head in bf16 and ForCausalLMLoss upcasts afterwards), and what the docstring for chunked_nll already claims ("same math as nll").

The same pattern is in the distillation trainers

DistillationTrainer._chunk pays it twice per chunk: once for the student and once for the teacher. Both are fixed here. The experimental async distillation trainer has a third copy, left alone: its projection runs outside accelerator.autocast(), so a bf16 GEMM there needs the weight cast explicitly, and that is not worth the risk in experimental code.

trl/trainer/utils.py:1446 looks similar and is deliberately not touched:

grad_hidden.add_(grad_logits @ w_chunk.float())

there grad_logits genuinely is fp32, so that upcast is important.

Only the SFT path is benchmarked below; the distillation change is mechanically identical and numerically free for the same reason, but I have not measured its speedup. Its test suite is green:

Numbers

One chunk, 256 tokens × vocab 248,320 × hidden 2048, 1×H100, bf16, fwd+bwd):

	fwd+bwd	peak mem
current (h.float() @ w.float().t())	23.37 ms	5.99 GB
this PR	3.86 ms (6.0×)	3.03 GB

In an 8×H100 profile of trl sft on Qwen3.6-35B-A3B (FSDP2, seq 4096, per-device batch 4, LoRA), the two fp32 SIMT GEMM kernels (sm80_xmma_gemm_f32f32_f32f32_f32_tn_n_...ffma, cutlass_80_simt_sgemm_...) account for 985 ms of a 4.57 s step (21.6% of all GPU kernel time) in 192 launches averaging 5.1 ms each.

End to end, on four released models, 16384 tokens per step, loss_type="chunked_nll" throughout:

model	mode	today	this PR
gemma-3-270m (vocab 262k)	full FT, 1×H100	24036	31609	1.32×
Qwen3-0.6B (vocab 152k)	full FT, 1×H100	21276	26641	1.25×
Qwen3-8B	full FT, 2×H100 FSDP2	3554	6009	1.69×
Qwen3-8B	LoRA r16, 2×H100 FSDP2	4531	7125	1.57×
Qwen3-30B-A3B (MoE)	LoRA attn, 2×H100 FSDP2	4251	5101	1.20×

tokens/s/GPU, at 16k tokens per step. The win grows with hidden_size, because that is what sets the fp32 GEMM's cost relative to the rest of the step. Which is why the

• 8B gains most and
• the MoE, whose 3B active parameters per token make the rest of the step cheap relative to its hidden=2048 projection, gains least.

Numerics

The projection now computes what nll computes, which is what the chunked_nll docstring already promises. Whether that shifts your loss curve depends on whether autocast covers the projection.

Under accelerate mixed precision it does not. _chunk runs inside autocast (the patch is applied before super().__init__(), so accelerate wraps the patched forward), and matmul is on autocast's lower-precision list, so autocast casts the explicitly .float()-ed operands straight back down to bf16. main is already on tensor cores there, and this PR only removes the wasted fp32 materialisation. One step at learning_rate=0.0, gradients captured before zero_grad and compared in fp32:

Qwen3-0.6B (vocab 152k)
  loss main=6.00852442  pr7=6.00852442
  params compared: 310
  global gradient relative difference: 0.00e+00
  worst per-parameter: 0.00e+00

gemma-3-270m (vocab 262k)
  loss main=5.02741528  pr7=5.02741528
  params compared: 236
  global gradient relative difference: 0.00e+00
  worst per-parameter: 0.00e+00

Bit-identical, so that run measures the memory saving and not the kernel change.

Without autocast over the projection (bf16 model, mixed_precision off) the GEMM does move from the fp32 SIMT path to bf16 tensor cores. That is the configuration the 21.6% profile above comes from, and there the numerics do change: logits round to bf16 before the softmax, and the backward grad_hidden / grad_weight matmuls run in bf16 rather than fp32. That is exactly what loss_type="nll" already does, but loss curves in that configuration will shift by bf16 rounding.

Why this matters for the cce PR

#6859
When measured against this "fixed" baseline instead of today's default, loss_type="cce" adds only 1.06× on Qwen3-8B (both full finetune and LoRA), against the 1.80×/1.66× it shows versus the unfixed path.
Nearly all of that apparent gain is this bug! cce still earns its place on models whose vocabulary is large relative to their size (1.76× on gemma-3-270m over the fixed baseline)

AsyncDistillationTrainer keeps the fp32 projection for now: its chunk runs outside accelerator.autocast(), so the dtype path there differs from these three and gets its own follow-up.

---

Note

Medium Risk
Changes the default chunked_nll training hot path and distillation JSD projection; numerics can shift when mixed precision does not autocast the matmul, though behavior is aligned with nll.

Overview
Chunked cross-entropy and distillation losses no longer upcast hidden states and lm_head weights to fp32 before the vocabulary GEMM. Projection is (h @ w.to(h.dtype).t()).float(), matching loss_type="nll" and Transformers’ causal LM loss: matmul runs in the hidden-state dtype (e.g. bf16 on tensor cores), then logits are fp32 only for softmax/CE.

On FSDP2, after full_tensor() on sharded lm_head weights, distillation now casts gathered fp32 master weights to the hidden dtype once instead of per-chunk fp32 copies.

Tests add test_bf16_hidden_fp32_weight for SFT chunked CE and distillation chunked JSD to cover bf16 activations with fp32 lm_head weights.

^{Reviewed by Cursor Bugbot for commit 6d4eab6. Bugbot is set up for automated code reviews on this repo. Configure here.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 00c9b2d569

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[albertvillanova]

The change itself looks right to me and should land. The strongest argument for it is actually not the one the description makes: with a bf16 model, nll / ForCausalLMLoss projects in bf16 and upcasts afterwards, so this PR makes the documented "same math as nll" claim (sft_config.py, sft_trainer.md, reducing_memory_usage.md) true, where today's fp32 upcast quietly diverges from it. Two things I would push back on before merging.

1. The numerical-equivalence claim is not supported by the experiment shown

global gradient relative difference: 0.00e+00 over 310 parameters means the two runs were bit-identical. As far as I can tell that is only possible in two configurations:

• full fp32 (h fp32, w fp32), where h.float() and w.to(h.dtype) are both no-ops;
• accelerate mixed precision with autocast active. _chunk runs inside autocast (the patch is applied before super().__init__() precisely so accelerate wraps the patched forward), and mm / matmul are on autocast's lower-precision list, so autocast casts the explicitly .float()-ed operands straight back down to bf16. main is already on tensor cores there, and the only thing this PR removes is the wasted materialization.

In both cases the patch is a no-op by construction, so that run cannot speak to the configuration the speedup comes from. Conversely, the two fp32 SIMT kernels at 21.6% of step time mean that profile had no autocast over the projection (model loaded in bf16, mixed_precision off), and there the numerics do change: logits are rounded to bf16 before the softmax, and the backward grad_hidden / grad_weight matmuls move from fp32 to bf16. Same in the FSDP2 mixed-precision case, where an fp32 master lm_head weight is now down-cast to bf16 for the GEMM, so the head's gradient is computed in bf16 rather than fp32.

None of that is an objection to the change, it is exactly what nll already does. But the two experiment blocks in the description describe mutually exclusive configurations, so "Numerically it changes nothing / this is not a speed-for-accuracy trade" overstates what was measured. I would restate it as "numerically identical to nll, which is what the docstring already promises", and add that bf16-weight runs will see loss curves shift by bf16 rounding.

2. Nothing tests the changed expression

Every test that reaches these lines pins dtype=torch.float32: TestChunkedCrossEntropyLoss._inputs, TestPatchChunkedCELMHead._setup / _setup_vlm, and the MoE variants. The only bf16 mentions in tests/test_sft_trainer.py are an xfail (test_train_padding_free) and a skipped gated slow test. So the green suite exercises (h @ w.to(h.dtype).t()).float() only where it reduces to the old code, and the mixed-dtype regression Codex caught above has nothing guarding it. A CPU-only unit test would close that:

def test_bf16_hidden_fp32_weight(self):
    """`h` bf16 against an fp32 master `lm_head` weight (FSDP2 mixed precision) must not raise."""
    hidden, weight, labels = self._inputs()
    loss_c, *_ = _chunked_cross_entropy_loss(hidden.bfloat16(), weight, self.CHUNK_SIZE, labels)
    loss_r, *_ = self._reference(hidden.bfloat16().float(), weight, labels)
    torch.testing.assert_close(loss_c, loss_r, atol=2e-2, rtol=2e-2)

Worth one for DistillationTrainer._chunk as well, which now has the same failure mode twice per chunk (student and teacher, and the teacher can be loaded in a different dtype than the student).

3. Comment consistency

Identical logic in the three copies now carries three different comments ("upcast only for the softmax, like ...", "upcast only afterwards, as ...", "as the other chunked projections do"). Per AGENTS.md, duplicated blocks should keep the same comment word for word, so I would align all three on one wording.

The async one is also inaccurate for its own path: it attributes the fp32 weight to FSDP2 mixed precision, but in the async trainer the projection runs outside self.accelerator.autocast() (async_distillation_trainer.py:1250), so plain DDP with bf16 already gives bf16 hidden states against an fp32 weight. That is the case Codex hit. The FSDP2 attribution is correct for SFT and distillation, where _chunk does run under autocast and only the gathered full_tensor() master weight can mismatch.

4. Minor: the cast is not free in the fp32-master-weight config

Where w is already fp32, the old .float() was a true no-op (returns self). The new .to(h.dtype) materializes a bf16 vocab-sized copy per chunk, and again on every gradient-checkpoint recompute. Still a clear net win, but it is the opposite of the memory story in the description and worth a sentence. In the async trainer it is easy to avoid: full_tensor() is already hoisted out of the chunk loop for exactly this reason (:1261), so the cast can be hoisted next to it. In SFT / distillation it cannot be hoisted naively, since under ZeRO-3 the weight has numel 0 outside maybe_gather_lm_head_ctx and a hoisted .to() would silently produce an empty tensor. Optional follow-up, not a blocker.

Checked and clean

• All four fp32 projections in the repo are covered, nothing missed.
• Leaving utils.py:1446-1447 alone is right: grad_logits there really is fp32, built from fp32 probability buffers, so those upcasts are load-bearing.
• No docs change needed, the "same math as nll" wording becomes more accurate, not less.
• The cast direction (weight to activation dtype) matches what lm_head's own forward computes in under both autocast and FSDP2 mixed precision, so no configuration regresses relative to loss_type="nll".

[qgallouedec]

@albertvillanova thks, all 4 addressed.

One correction: the FSDP2 attribution is backwards (probed on a 2xH100 FSDP2 bf16 run)

• SFT projects inside the patched lm_head.forward, so FSDP2 has already gathered and cast the parameter: Parameter dtype=bfloat16. The cast is a no-op there.
• The divergence loss projects outside any forward: DTensor dtype=float32, and full_tensor() returns fp32. Nothing casts it back either, Trainer.autocast_smart_context_manager is a nullcontext.

1. description rewritten. The bit-identical run was under autocast, so it measures the memory saving, not the kernel. Without autocast the numerics do shift by bf16 rounding, same as nll.
2. tests added for both paths.
3. comments aligned. Async reverted entirely: its projection runs outside autocast and it's experimental, not worth the risk here, low prority in follow up pr
4. cast hoisted next to the full_tensor(), once per step instead of once per chunk. The per-chunk one stays for ZeRO-3, where the weight is only gathered inside maybe_gather_lm_head_ctx.

[albertvillanova]

Re-reviewed at be75407. All four are addressed, and the rewritten "Numerics" section is exactly the claim the evidence supports. Thanks for probing the FSDP2 side rather than taking my word for it, I had the two paths the wrong way round.

LGTM.

Just a couple of comments, both comment-level, neither blocking:

1. Which head did you probe in the divergence loss? _ForwardRedirection calls wrapper_module(*args, **kwargs) (trl/models/utils.py:370), i.e. __call__, so the student root's FSDP2 pre-forward hooks do fire and student_lm_head.weight should be gathered and cast by the time _compute_loss runs, same as SFT. The teacher is the one I would expect to stay a sharded fp32 DTensor: its FSDP group is not in forward while the student's redirection is running, only its backbone got its own redirection earlier (:1905). If that is what you saw, then the student half of the new cast is a no-op in that config and only the teacher half is load-bearing. Either way the code is right, this is only about what the comment implies.

2. Now sft_trainer.py:321-328 is the odd one out. Same full_tensor() block, no cast, and the comment still asserts "Under FSDP2, lm_head.weight is a DTensor". That is fine, but it needs a comment saying why, otherwise the next reader will "restore consistency" by adding one.

[qgallouedec]

thanks, added a comment for both