From faee7d2321bc093942ec687984f438538b6d5f39 Mon Sep 17 00:00:00 2001 From: Vaibhav Jindal Date: Sun, 2 Aug 2026 01:55:39 +0000 Subject: [PATCH] fix(ce): restore true-class gradient precision in low-precision buffers #1268 replaced the per-element `tl.where(X_offsets != y, ...)` selection with a post-loop read-modify-write of the true-class gradient: tl.store(X_ptr + y, tl.load(X_ptr + y) + dxy) X_ptr is the logits buffer, which is bf16/fp16 in mixed-precision training (FLCE computes the gradient in place). The round-trip therefore rounds softmax(x_y) / N to the buffer dtype *before* the -(1 - label_smoothing) / N term is subtracted. Since dx_y = (softmax(x_y) - 1) / N cancels catastrophically as softmax(x_y) -> 1, the pre-rounding error is amplified by roughly 1 / (1 - softmax(x_y)). Measured max relative error of dx_y against an fp64 reference (bf16, V=32000, H100): softmax(x_y) before #1268 at #1268 torch 0.90 0.24% 1.95% 0.43% 0.99 0.30% 19.99% 0.46% 0.999 0.38% 100.00% 0.42% Confident predictions are the common case once a model has trained for a while, so this silently degrades the dominant gradient signal. It surfaced as a convergence failure in test_mini_model[mini_gemma3_text-32-1e-05-dtype25-0.01-0.01-0.1-0.01-0.01-0.01], where 32 steps of accumulated drift push a top-k logprob past tolerance. Recompute dx_y once in fp32 after the loop, from the already-live ori_X_y, m and d, and write it with a single store. This keeps #1268's O(1) correction (no O(V) compare/select is reintroduced) while folding in the -(1 - label_smoothing) term before the result is rounded to the buffer dtype. Accuracy returns to the pre-#1268 values exactly (0.02% / 0.36% / 0.24% / 0.30% / 0.38%), and the speedup is retained: 3.585 ms vs 3.571 ms (this PR vs #1268) at BT=8192, V=128256 bf16 forward+backward on H100, against 3.619 ms before #1268. The existing CE suite passed with the regression because its random logits leave softmax(x_y) near zero, so add a regression test that drives softmax(x_y) -> 1 and requires the true-class gradient to be no less accurate than torch's own low-precision backward. It fails on the #1268 kernel and passes here with ~6x headroom. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/liger_kernel/ops/cross_entropy.py | 37 +++++++++++----- test/transformers/test_cross_entropy.py | 57 +++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 11 deletions(-) diff --git a/src/liger_kernel/ops/cross_entropy.py b/src/liger_kernel/ops/cross_entropy.py index 08fa6c281..6063bec91 100644 --- a/src/liger_kernel/ops/cross_entropy.py +++ b/src/liger_kernel/ops/cross_entropy.py @@ -245,21 +245,36 @@ def liger_cross_entropy_kernel( tl.store(X_ptr + X_offsets, X_block, mask=X_offsets < n_cols) - # dx_y correction: apply the -(1 - label_smoothing) term once at index y - # (replaces the per-element tl.where removed above). Barrier first so the loop's - # in-place store to X[y] is visible before we read it back. + # dx_y correction: recompute the true-class gradient once, in fp32, and overwrite X[y] + # (replaces the per-element tl.where removed above). The -(1 - label_smoothing) term must + # be folded in *before* the result is rounded to X's dtype: dx_y = (softmax(x_y) - 1) / N + # cancels catastrophically as softmax(x_y) -> 1, so a read-modify-write of the value the + # loop already stored would lose most of the significant bits in bf16/fp16. + # ori_X_y is the (softcapped) fp32 logit at y, so softmax_X_y matches the loop exactly. + # Barrier first so the loop's in-place store to X[y] cannot land after this store. tl.debug_barrier() - dxy = -(1 - label_smoothing) - if HAS_WEIGHT: - dxy = dxy * weight_y + softmax_X_y = tl.exp2((ori_X_y - m) * LOG2_E) / d + if not HAS_WEIGHT: + dx_y = softmax_X_y + dx_y += 2 * lse_square_scale * lse * dx_y + dx_y += -eps + dx_y += -(1 - label_smoothing) + if reduction == "mean": + dx_y = dx_y / n_non_ignore + else: + dloss_ori_y = (1 - label_smoothing) * softmax_X_y - (1 - label_smoothing) + dloss_ori_y = dloss_ori_y * weight_y + dloss_smooth_y = eps * (-weight_y + softmax_X_y * weight_sum) + dz_loss_y = 2 * lse_square_scale * lse * softmax_X_y if reduction == "mean": - dxy = dxy / sum_non_ignore_weight - elif reduction == "mean": - dxy = dxy / n_non_ignore + dloss_ori_y = dloss_ori_y / sum_non_ignore_weight + dloss_smooth_y = dloss_smooth_y / sum_non_ignore_weight + dz_loss_y = dz_loss_y / n_non_ignore + dx_y = dloss_ori_y + dloss_smooth_y + dz_loss_y if HAS_SOFTCAPPING: t_y = ori_X_y / softcap - dxy = dxy * (1 - t_y * t_y) - tl.store(X_ptr + y, tl.load(X_ptr + y) + dxy) + dx_y = dx_y * (1 - t_y * t_y) + tl.store(X_ptr + y, dx_y) # We need tl.debug_barrier() to ensure the new result of X_ptr is written as mentioned in # https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/ops/cross_entropy.py#L34 diff --git a/test/transformers/test_cross_entropy.py b/test/transformers/test_cross_entropy.py index 38f4071de..369cb7f81 100644 --- a/test/transformers/test_cross_entropy.py +++ b/test/transformers/test_cross_entropy.py @@ -1265,3 +1265,60 @@ def test_correctness_with_predicted_tokens(B, T, V, ignore_index, dtype): # Verify backward still works result.loss.backward() assert _input.grad is not None + + +@pytest.mark.parametrize( + "dtype", + [ + pytest.param( + torch.bfloat16, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + pytest.param(torch.float16), + ], +) +@pytest.mark.parametrize("label_smoothing", [0.0, 0.1]) +@pytest.mark.parametrize("reduction", ["sum", "mean"]) +def test_correctness_true_class_grad_confident_predictions(dtype, label_smoothing, reduction): + """ + dx_y = (softmax(x_y) - (1 - label_smoothing)) / N cancels catastrophically once the model is + confident, so the -(1 - label_smoothing) term must be folded in before the result is rounded + to the low-precision dtype of the in-place gradient buffer. Random logits leave softmax(x_y) + near zero and never expose a lossy round-trip through that buffer; confident ones do. + + The true-class gradient is checked against an fp32 reference computed from the same + low-precision logits, and is required to be no less accurate than torch's own low-precision + backward on those logits. + """ + torch.manual_seed(0) + B, T, V = 2, 64, 4096 + + logits = torch.randn(B * T, V, device=device, dtype=torch.float32) + target = torch.randint(0, V, (B * T,), device=device, dtype=torch.long) + # Drive softmax(x_y) close to 1 so the true-class gradient is a small difference of + # comparatively large terms. + logits[torch.arange(B * T, device=device), target] = logits.max(dim=-1).values + 10.0 + _tensor = logits.to(dtype) + + _input = _tensor.detach().clone().requires_grad_(True) + _input2 = _tensor.detach().clone().requires_grad_(True) + _input_ref = _tensor.detach().clone().float().requires_grad_(True) + + torch_ce = CrossEntropyLoss(reduction=reduction, label_smoothing=label_smoothing) + liger_ce = LigerCrossEntropyLoss(reduction=reduction, label_smoothing=label_smoothing) + + torch_ce(_input, target).backward() + liger_ce(_input2, target).backward() + torch_ce(_input_ref, target).backward() + + rows = torch.arange(B * T, device=device) + ref = _input_ref.grad[rows, target].double() + torch_err = (_input.grad[rows, target].double() - ref).abs().max().item() + liger_err = (_input2.grad[rows, target].double() - ref).abs().max().item() + + # Rounding the softmax term before the subtraction inflates this error by ~1/(1 - softmax(x_y)), + # i.e. orders of magnitude, so a small constant factor over torch is a wide margin. + assert liger_err <= max(4 * torch_err, torch.finfo(dtype).tiny), ( + f"true-class grad error {liger_err:.3e} exceeds torch's {torch_err:.3e} " + f"(reduction={reduction}, label_smoothing={label_smoothing}, dtype={dtype})" + )