Add loss_type="cce", a fused linear cross-entropy loaded from the Hub

[qgallouedec]

Adds SFTConfig(loss_type="cce"): a fused linear cross-entropy that computes the loss of hidden @ lm_head.T without ever materialising the (tokens, vocab) logit matrix. The kernel is fetched through kernels, so this adds no pip dependency.

SFTConfig(loss_type="cce")

Benchmark

Stacked on #6863, so the baseline is chunked_nll with the projection fix (against today's chunked_nll, most of the apparent gain would really be the fp32-projection bug #6863 fixes). Same GPU and config per row, 1×H100 unless noted, seq 2048:

model	mode	vs chunked_nll + #6863	peak VRAM
gemma-3-270m (vocab 262k)	full FT	1.65x	3.1 → 2.6 GB
Llama-3.2-1B (vocab 128k)	full FT	1.33x	10.3 → 10.3 GB
Qwen3-0.6B (vocab 152k)	full FT	1.32x	5.2 → 5.2 GB
Qwen3-8B	full FT	1.07x	61.1 → 61.1 GB
Qwen3-8B	LoRA r16, 8k tok/step	1.06x	19.7 → 19.7 GB
Qwen3-8B	LoRA r16, 2k tok/step	1.02x	16.7 → 16.6 GB
Qwen3-30B-A3B (MoE)	LoRA r16, FSDP2×EP, 4×H100	1.02x	19.1 → 19.1 GB

This is a throughput change, not a memory one: chunking already bounds the loss memory, and the peak is unchanged at every size measured (on gemma-3-270m at 64k and 128k tokens per step the two are identical to the decimal, while cce is 1.61x faster).

When it is worth turning on

The gain is the share of the step spent in the lm_head projection + cross-entropy, which tracks vocabulary size relative to the rest of the model. Measured against the fixed baseline at 16k tokens per step:

model	vocab / (layers × hidden)	speedup
Qwen3-8B	1.03	1.06x
Qwen3-0.6B	5.30	1.29x
gemma-3-1b	8.75	1.39x
gemma-3-270m	22.76	1.76x

Rule of thumb: worth it above a ratio of roughly 5, marginal near 1.

What it supports

Everything chunked_nll does: an lm_head bias, a non-unit logit_scale (Cohere / Command-R), float32 hidden states, and final_logit_softcapping. mean_token_accuracy and entropy come out of the same fused pass and are logged as usual — unlike use_liger_kernel=True, which drops them. _cut_cross_entropy_loss takes the same arguments and returns the same 4-tuple as _chunked_cross_entropy_loss.

Correctness

Against chunked_nll on Qwen3-0.6B, one step at learning_rate=0.0:

	chunked_nll	cce	rel diff
loss	6.008524	6.006553	3.3e-04
entropy	4.381336	4.385639	9.8e-04

mean_token_accuracy differed by one position on that run (a bf16 near-tie argmax flip); with targets set to the true fp32 argmax the kernel reports exact accuracy at V = 200000, 32000 and 150000.

Before merging

• The kernel is at qgallouedec/fused-linear-ce and should move to trl-lib first — a TRL default should not point at a personal namespace.
• trust_remote_code=True is needed until it lives under an organisation with trustedKernelPublisher enabled on the Hub.

Both are marked TODO in the source.

---

Note

Medium Risk
Opt-in change to the core SFT loss path that loads and executes a pinned remote Hub kernel with trust_remote_code; distributed-training edge cases are explicitly handled and covered by new tests.

Overview
Adds SFTConfig(loss_type="cce") for Cut Cross-Entropy SFT: next-token loss is computed with the lm_head fused into a Hub-loaded kernel (trl-lib/fused-linear-ce, pinned revision) so the full [tokens × vocab] logits tensor is never materialized, targeting better throughput on large-vocabulary models while matching the chunked_nll training path (forward patch, MoE aux loss, mean_token_accuracy / entropy from the same fused pass).

cce reuses the existing chunked-CE forward patch via use_cce on _patch_chunked_ce_lm_head, with a new _cut_cross_entropy_loss helper (including an all--100 labels branch so DDP/FSDP backward still touches every parameter). kernels>=0.14.0 is added to the kernels optional extra; version gating and xfail_old_kernels tests enforce get_kernel(..., trust_remote_code=True). Docs add the CCE paper entry in paper_index.md. cce remains incompatible with use_liger_kernel=True (same as chunked_nll).

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

[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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b0a1675840

ℹ️ 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".

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 0aa3e50. Configure here.}

[albertvillanova]

Reviewed on top of #6863 (which I reviewed separately). Fetched both branches and checked the claims against the code, the Hub and the CI logs.

First, a few things I verified so nobody has to re-check them:

• The pin is valid. 00e29bdf4e142fafa8f558be414c7a9f344d61d0 is the head of the v1 branch of the kernel-type repo, via https://huggingface.co/api/kernels/trl-lib/fused-linear-ce, currently equal to main. Note for anyone verifying: the api/models/... namespace hosts a different repo of the same name with a different commit graph and returns 404 for this sha. I fell into that first.
• The kernel's math matches _chunk operation for operation: matmul in the compute dtype with fp32 accumulation, then + bias in fp32, then * logit_scale, then softcap. entropy_sum = lse - sum(p * logit) is the correct -sum(p log p), and the argmax is taken on the scaled/softcapped logits, so both metrics are consistent with the chunked path.
• The weight.to(hidden.dtype) cast is load-bearing: the kernel raises on a dtype mismatch.
• CI genuinely ran both new tests (kernels 0.16.1 resolves through transformers in the dev and latest jobs). They did not skip or xfail.

---

Blocking

1. trust_remote_code=True is hardcoded, and trl-lib is still not a trusted publisher

The Hub API reports "trustedPublisher": false on the kernel repo, so the flag is indeed still required. The consequence is that selecting loss_type="cce" downloads and executes remote Python from the Hub with no user opt-in, and independently of SFTConfig.trust_remote_code.

Of the two items in the "Before merging" section, this is the only one still outstanding. At minimum the SFTConfig docstring has to say so; today it says only "Requires kernels to be installed", which does not prepare anyone for remote code execution.

Correctness

2. cce does not drop ignored tokens before the projection, but chunked_nll does

This is the finding I would most like an answer on.

_chunked_cross_entropy_loss packs valid tokens to the front with argsort and runs ceil(n_valid / chunk_size) chunks, so its cost scales with n_valid. The kernel does not do this: _fwd_kernel and _bwd_d_kernel iterate over all N rows unconditionally, validity only zeroes the per-row weight w after the GEMM, and the backward torch.mm(d.t(), hidden) likewise spans every row. So cce's cost scales with total positions, not valid ones.

For completion-only or prompt-completion SFT (very common in TRL) at, say, 30% valid tokens, cce performs roughly 3x the projection work chunked_nll does. Measured against speedups of 1.02x to 1.65x on all-valid batches, that plausibly makes cce slower than the default it is being compared against in that regime. The benchmark table does not name the dataset, and every row reads like plain language modeling, so this case looks unmeasured.

The fix is cheap: gather before the call (hidden_states[valid], targets[valid]). The data-dependent output shape that rules this out for chunked_nll (XLA compilation) does not apply here, since this path is Triton/CUDA only. Either way, one benchmark row on a completion-only dataset would settle it.

3. The PEFT lm_head guard names the wrong loss type

trl/trainer/sft_trainer.py:1451 still hardcodes loss_type='chunked_nll', but the branch now serves cce too, so a cce user gets an error about a loss type they did not select. Use args.loss_type, exactly as the PR already does at line 1468.

4. The loss_type CLI help was not updated

trl/trainer/sft_config.py:284-289: the field's metadata["help"], which is what --loss_type prints, still lists only 'nll', 'dft' and 'chunked_nll'. Only the class docstring was updated.

5. The empty-batch branch re-introduces the fp32 weight copy that #6863 exists to delete

loss = (hidden_states.float().sum() + lm_head_weight.float().sum()) * 0.0

lm_head_weight.float() materializes a full (V, H) fp32 temporary before reducing it to a scalar. For a 262k-vocab head that is the same ~2 GB copy #6863 removes, and it fires precisely in the fully-masked micro-batch case. lm_head_weight.sum().float() produces the identical value without the copy, same for hidden_states.

Consistency (AGENTS.md)

6. The shift diverges from the chunked path

torch.roll(labels, -1) plus masking the last column, versus the chunked path's hidden_states[..., :-1, :] / labels[..., 1:] slice. Numerically equivalent, but the duplicated blocks are meant to stay aligned, and the slice form additionally drops the last row of each sequence that cce currently pushes through a full vocabulary sweep for nothing.

7. No labels / shift_labels guard

The chunked path raises "At least one of 'labels' or 'shift_labels' must be provided." and test_requires_labels_or_shift_labels covers it. _cut_cross_entropy_loss instead reaches torch.roll(None, ...) and raises a TypeError.

8. correct has two different dtypes

fp32 in the empty branch, int64 from the kernel otherwise. The chunked path returns fp32 in both cases.

9. _patch_chunked_ce_lm_head docstring

Still says the loss is computed "via [_chunked_cross_entropy_loss]", with no mention of the fused path. The new use_cce entry is also documented before is_vlm, while the signature has is_vlm first.

Test coverage

10. The only cce unit test never reaches the kernel

test_cce_all_ignored_returns_zero hits the n_valid == 0 early return, so it exercises pure PyTorch and nothing else. TestChunkedCrossEntropyLoss has nine numerical tests for the chunked path and zero for cce; the only thing that touches the kernel is test_train_cce_loss, which asserts "not None" and "params changed".

That leaves the numerical evidence in the description (3.3e-04 on loss, 9.8e-04 on entropy) as a one-off manual measurement nothing will re-run. logit_scale and final_logit_softcapping are listed under "What it supports" and are entirely untested on this path.

Mirroring test_forward_matches_cross_entropy, test_backward_matches_reference, test_lm_head_bias, test_num_items_in_batch_reduction and test_shift_labels_matches_labels against _cut_cross_entropy_loss with a loose tolerance would cover it.

11. The strict xfail is load-bearing on an unused call

test_cce_all_ignored_returns_zero only xfails on an old kernels because _fused_linear_cross_entropy() is called at the top of the function, even though the empty branch never uses its result. Moving that call below the early return, which is a natural cleanup, silently turns the test into a strict-xfail XPASS failure.

Documentation

12. "Peak memory does not scale with vocab_size" is misleading next to the default

Both the SFTConfig docstring and the paper-index entry lead with it. It is true against "nll", but the default is "chunked_nll", which already bounds that memory, and the PR's own table states peak VRAM is unchanged at every size measured. Sitting directly beside the chunked_nll bullet, it reads as a memory win that does not exist. The description gets this right ("This is a throughput change, not a memory one"); the user-facing docs should say the same.

13. Missing caveats in the cce bullet

It omits the use_liger_kernel incompatibility that the code enforces (the chunked_nll bullet has it), and omits that the loss type downloads and executes a kernel from the Hub.

14. Wording

"the lm_head projection is fused into the cross-entropy kernel by a fused hub kernel" is redundant, and "hub" should be "Hub".

15. speeding_up_training.md and kernels_hub.md are untouched

Both are natural homes for a throughput feature backed by a Hub kernel. In particular the "when it is worth turning on" rule of thumb (vocab / (layers x hidden) above roughly 5) is the most actionable thing in this PR and currently lives only in the description.

Nits

• trl/trainer/sft_trainer.py:1437 is 125 characters against the file's 119. Ruff's select does not include E501, so CI will not flag it.
• The "Before merging" section is stale: the kernel already moved to trl-lib, and git grep TODO on the head is empty, so "Both are marked TODO in the source" no longer holds. Only the trust_remote_code item still stands.
• tests/testing_utils.py imports transformers' is_kernels_available as is_kernels_installed right beside TRL's own is_kernels_available, which means something different (version-gated). Naming TRL's helper after what it actually gates would avoid the trap.

Optional

tests/invariant/ already has an sft equivalence class. A cce config there would give the correctness claim a snapshot that actually re-runs on every version bump. It runs fp32, where the kernel takes the IEEE path, so agreement should be well inside the harness's absolute tolerance.

[qgallouedec]

after second thought, I think we should actually close this one:
The win is real only for small models with a huge vocab, where the loss is actually the global peak. But for any medium size model, the peak will be for the activations. And the speed up tiny compared to the rest of the forward.

Not worth a permanent loss_type I think. Maybe revisit in the future, but for now the balance isn't favorable.