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. And the experimental async distillation trainer has a third copy. All four are fixed here.

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.

Numerically it changes nothing

This is not a speed-for-accuracy trade. One step at learning_rate=0.0 on Qwen3-0.6B, 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

That is expected once stated: h and w are already bf16, so upcasting them is lossless, and a bf16 tensor-core GEMM accumulates in fp32 regardless.

Both paths compute the same products at the same accumulation precision, the fp32 version just does it on
the SIMT path and materialises an fp32 copy of the lm_head weight to get there.

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)

---

Note

Low Risk
Localized numeric-path change in loss helpers; PR reports identical loss and zero gradient drift when operands are already bf16, with comments documenting the FSDP2 dtype edge case.

Overview
Chunked lm_head projections in SFT (loss_type="chunked_nll"), DistillationTrainer, and async distillation no longer upcast hidden states and weights to fp32 before the matmul. They now compute (hidden @ weight.to(hidden.dtype).T) and only then cast logits to fp32 for softmax/NLL or JSD—aligned with "nll", ForCausalLMLoss, and the intended “same math as nll” behavior.

This keeps the large vocabulary GEMM on tensor cores (bf16/fp16) instead of fp32 SIMT, avoiding repeated fp32 copies of the full lm_head weight each chunk (and on gradient-checkpoint recomputation). Weights are explicitly cast to the hidden-state dtype so FSDP2 mixed precision still matches what lm_head.forward would do when master weights are fp32.

^{Reviewed by Cursor Bugbot for commit f99ca55. 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.