[JobQiu]

Summary

Two changes to the standalone cross entropy (LigerCrossEntropyLoss / LigerCrossEntropyFunction).

1. Per-token loss precision. Standalone CE stores the per-token loss in the input dtype (cross_entropy.py:379 on main), so bf16 per-token losses carry bf16 rounding (FLCE has kept fp32 since 2024-08). Per-token loss and z_loss are now always computed and reduced in fp32 internally, and mean/sum are cast once to the input dtype (output dtype unchanged). A new loss_dtype=torch.float32 option returns the fp32 loss; it is supported by the Triton backend only, other backends raise ValueError when it is set. With the default bf16 output the returned loss error is bounded by the final bf16 cast, so the default path shows no measurable difference in the benchmark below; the precision gain is for callers who set loss_dtype=torch.float32.

2. reduction="none" backward memory, as an opt-in. By default (in_place=False) the backward is exactly main's code path: it multiplies the saved gradient out of place by the per-token upstream gradient, so retain_graph=True and backward through different slices of the per-token losses keep working. The new in_place=True option (Triton backend only; other backends raise ValueError when it is set) uses a row-wise in-place Triton kernel that scales the saved gradient by grad_output[row] for any 1-D grad (mean, sum, weighted, strided), with an fp32 multiply and one rounding. It avoids the extra BT x V temporary and has no device-to-host sync, and gradients are bitwise identical to main. With in_place=True the graph is single-use: a second backward raises RuntimeError. reduction="mean"/"sum" are unchanged from main, except that a second backward after a 0-D upstream gradient other than 1.0 now raises RuntimeError instead of silently returning a wrong gradient (with an upstream gradient of exactly 1.0, retain_graph=True still works).

Related: #678 / #680 (introduced the per-row multiply), #1292 and #1251 (open; textual conflicts only, and this PR tests grad_output.ndim > 0 first to match #1292's layout), and the FLCE precedent.

Details

The benchmark script gains --reduction {mean,sum,none_sum,none_mean}, --dtype and --in-place options (defaults unchanged; --in-place requires a none_* reduction and otherwise exits with an error). The backward speed mode now times each backward on a fresh graph: the forward is rerun untimed, _input.grad is reset and the 256 MB L2 cache is flushed before the timed region, and the loop is sized like triton.testing.do_bench (one discarded call, a 5-call average estimate, 25 ms warmup, 10 ms rep, quantiles 0.5/0.2/0.8). This fixes a RuntimeError the benchmark hit when a backward was repeated on the same graph. The forward, full and no-grad-forward modes and the memory benchmark are unchanged (they already build a fresh graph per call).

Benchmark: 1x NVIDIA H100 80GB HBM3, torch 2.13.0.2+cu130, triton 3.7.1, BT=8192, bf16 random logits x4, targets random with every 17th ignored, 20 iterations after warmup, forward + backward on a fresh graph per iteration, public vocab sizes. "default" is in_place=False; extra peak memory is in GiB.

vocab	usage	main ms	default ms	in_place ms	main peak	default peak	in_place peak
32000	none+mean	1.625	1.599	1.299	1.953	1.953	0.977
32000	none+sum	1.577	1.543	1.307	1.953	1.953	0.977
32000	none+weighted	1.630	1.605	1.298	1.953	1.953	0.977
32000	mean	1.062	1.093	1.057 (no effect)	0.977	0.977	0.977
128256	none+mean	6.046	6.020	4.797	7.828	7.828	3.914
128256	none+sum	5.863	5.827	4.808	7.828	7.828	3.914
128256	none+weighted	6.026	5.974	4.803	7.828	7.828	3.914
128256	mean	3.520	3.501	3.488 (no effect)	3.914	3.914	3.914
151936	none+mean	7.293	7.198	5.799	9.273	9.273	4.637
151936	none+sum	7.071	7.036	5.809	9.273	9.273	4.637
151936	none+weighted	7.267	7.220	5.807	9.273	9.273	4.637
151936	mean	4.316	4.232	4.244 (no effect)	4.637	4.637	4.637

Gradients are bitwise equal to main for both paths in all 24 combinations (12 cases x {default, in_place=True}). On the default path, the retain_graph slice case (loss[:k].sum().backward(retain_graph=True) then loss[k:].sum().backward(), BT=8192, k=BT/3) succeeds and the accumulated gradient equals a single backward bitwise, for all three vocab sizes.

The B300 tables below were measured on the previous head (ddd3ec4), on 1x NVIDIA B300 SXM6, torch 2.13.0, same setup as above. Their "PR" columns are the in_place=True path: the forward kernel and the in-place kernel are unchanged since. The default path was not re-measured on B300 because it is main's code path. Extra peak memory is main -> PR.

vocab	usage	main ms	PR ms	speedup	extra peak memory (GiB)
32000	none+mean	1.07	0.65	1.65x	1.95 -> 0.98
32000	none+sum	1.03	0.64	1.60x	1.95 -> 0.98
32000	none+weighted	1.07	0.65	1.65x	1.95 -> 0.98
32000	mean	0.55	0.56	0.99x	0.98 -> 0.98
128256	none+mean	3.91	2.23	1.76x	7.83 -> 3.91
128256	none+sum	3.74	2.22	1.68x	7.83 -> 3.91
128256	none+weighted	3.90	2.25	1.74x	7.83 -> 3.91
128256	mean	1.69	1.71	0.99x	3.91 -> 3.91
151936	none+mean	4.72	2.70	1.75x	9.27 -> 4.64
151936	none+sum	4.49	2.70	1.66x	9.27 -> 4.64
151936	none+weighted	4.69	2.70	1.73x	9.27 -> 4.64
151936	mean	2.08	2.09	0.99x	4.64 -> 4.64

Per-token loss precision (forward unchanged since the previous head, so measured there). On B300, gradients are bitwise equal to main in all 12 cases, and the per-token loss max abs error vs an fp32 PyTorch reference is:

vocab	default (bf16 output)	loss_dtype=torch.float32
32000	0.0625	1.9e-6
128256	0.12	3.8e-6
151936	0.117	3.8e-6

On H100 (also on the previous head) the per-token loss max abs error is 0.093-0.124 with the default bf16 output and 3.8e-6 with loss_dtype=torch.float32.

The code, tests, docs and this description were written with AI assistance (GitHub Copilot). I reviewed every changed line; the tests and benchmarks listed above were run on the hardware described.

Testing Done

On 1x NVIDIA H100 80GB HBM3 (torch 2.13.0.2+cu130, triton 3.7.1), head a148e0c:

• test/transformers/test_cross_entropy.py plus test/cutedsl/test_cross_entropy.py on this branch: 480 passed, 1 skipped.

• This branch's test file run against main's code: 121 failed / 191 passed, and the failures are exactly the tests that need the change: the in_place kwarg tests (per-row multiply 48, matches torch 24, no-BT-x-V-temporary memory 12, loss_dtype keeps input dtype 2, retain_graph slices 2, in_place second backward 1, in_place backend check 1), the loss_dtype tests (24 + 4 + 1 + 1), and the scalar second-backward guard test (1).

• Repo benchmark script (--sweep-mode model_config --bt 4096): on the previous head the reported failure reproduces (RuntimeError in the backward speed mode with the default arguments, exit 1). On this head the script completes (exit 0) with the default arguments, with --reduction none_sum, and with --reduction none_sum --in-place; --in-place with the default mean reduction exits with --in-place requires --reduction none_sum or none_mean. These runs shared the GPU with the full test suite, so their timings are not quoted.

• Full suite (pytest test/ --ignore=test/convergence) on this head: 2 failed, 6172 passed, 1558 skipped, 21 xfailed, 1 error. The failing set is exactly the pre-existing one seen on main in the previous H100 round (test_rope.py::test_qwen3_5_vision_rope_matches_hf, and test_low_probability_vocab_block_is_finite in test_jsd.py and test_fused_linear_jsd.py); the 41 extra passes are the new cross-entropy tests.

• ruff check and ruff format --check pass on every touched file.

Convergence on the B300 (previous round, pytest test/convergence): the run hits CUDA error: unspecified launch failure after the first 7 mini-model tests, and every later test errors at setup. main and the branch behave identically (7 passed, 16 skipped, 12 xfailed, 310 errors with -k "not gpt_oss"; same first failing test), so it is an environment issue on that machine, not this change. It would be good to see the convergence job in CI.

• Hardware Type: NVIDIA H100 80GB HBM3 (this round); the previous round also used NVIDIA B300 SXM6
• run make test to ensure correctness (CE tests, see above)
• run make checkstyle to ensure code style (touched files)
• run make test-convergence to ensure convergence (blocked by an environment crash that also happens on main, see above)

[vaibhavjindal]

@JobQiu
The fp32 precision and backward-memory improvements look worthwhile, but I'd request changes before merging:

1.The new guard breaks benchmark_cross_entropy.py: its backward benchmark repeatedly calls backward(..., retain_graph=True) on the same graph, so it now raises on the second iteration--even with the default reduction. Please update the benchmark accordingly.
2. Also, although documented, removing retained-graph support for reduction="none" breaks previously working first-order usage, such as backpropagating separately through different slices of the per-token losses. Please preserve a non-destructive path or make destructive backward opt-in. Simply removing the guard would not be safe because the saved gradients have already been modified.

[JobQiu]

Thanks for the review, both points were valid. Fixed in a148e0c.

1. Reproduced on the previous head: the backward speed benchmark raised on the second y.backward(do, retain_graph=True) even with the default reduction, because do = torch.randn_like(y) is not 1.0, so the scalar in-place path ran and the guard fired. The CE benchmark now times the backward mode on a fresh graph per call (forward rerun untimed, grad reset and L2 flush before the timed region, do_bench-style loop sizing); the other modes already build a fresh graph per call. It also gained --in-place for the none_* reductions. On this head the script completes with the default arguments, with --reduction none_sum, and with --reduction none_sum --in-place.

2. The destructive path is now opt-in. in_place: bool = False was added to LigerCrossEntropyLoss, liger_cross_entropy and LigerCrossEntropyFunction (Triton backend only; the other backends raise ValueError on True). By default reduction="none" uses the same out-of-place per-row multiply as main, so retain_graph=True and backward through slices of the per-token losses keep working. test_none_reduction_retain_graph_slices_match_single_backward checks that backprop through loss[:k] then loss[k:] equals a single backward bitwise (and matches torch in fp32), and that in_place=True raises on the second backward. With in_place=True the row-wise in-place kernel and the guard are kept. For mean/sum nothing changed in how the gradient is scaled; the only difference from main is that a second backward after a 0-D upstream gradient other than 1.0 now raises instead of returning a wrong gradient, and with an upstream gradient of 1.0 retain_graph=True still works.

The description is updated with the re-run on H100 (default path matches main in time and peak memory, in_place=True keeps the earlier gains, gradients bitwise identical to main for both paths in all 24 combinations) and the docs now describe the default and the opt-in.

A possible follow-up (not in this PR): when the reduction="none" upstream gradient is all ones (loss.sum().backward()), the multiply can be skipped entirely like the existing 0-D 1.0 path, at the cost of one torch.equal sync. Measured on H100 it is faster than in_place=True with the same peak memory and keeps the graph reusable. Happy to add it here instead if you prefer.