fix(ce): restore true-class gradient precision in low-precision buffers

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

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})"
+    )