fix: torch.compile loss corruption at ignore_index rows (+ perf: drop BT x V loss buffer in FLJSD)

src/liger_kernel/ops/cross_entropy.py

@@ -106,6 +106,12 @@ def liger_cross_entropy_kernel(
         for i in range(0, n_cols, BLOCK_SIZE):
             X_offsets = i + tl.arange(0, BLOCK_SIZE)
             tl.store(X_ptr + X_offsets, 0.0, mask=X_offsets < n_cols)
+        # Zero the losses explicitly rather than relying on the caller having zeroed them: under
+        # `torch.compile` the buffer handed to us is a functionalization clone, so caller-side
+        # initialization is not guaranteed to reach it.
+        tl.store(loss_ptr + program_id * loss_stride, 0.0)
+        if RETURN_Z_LOSS:
+            tl.store(z_loss_ptr + program_id * loss_stride, 0.0)
         # For ignored tokens, set token accuracy to 0
         if RETURN_TOKEN_ACCURACY:
             token_accuracy_ptr += program_id * token_accuracy_stride

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -150,10 +150,21 @@ def fused_linear_cross_entropy_forward(
             scaling_factors = pred_probs.detach()  # Detach to ensure no gradient flow
 
         # unreduced loss
-        loss_1d_slice = loss_1d[start_idx:end_idx]  # chunk_size,
-        z_loss_1d_slice = z_loss_1d[start_idx:end_idx] if return_z_loss else None
-        token_accuracy_1d_slice = token_accuracy_1d[start_idx:end_idx] if return_token_accuracy else None
-        predicted_tokens_1d_slice = predicted_tokens_1d[start_idx:end_idx] if return_predicted_tokens else None
+        # NOTE: these must be standalone contiguous buffers, *not* `xxx_1d[start_idx:end_idx]`
+        # views. `torch.compile` functionalizes the kernel's mutation by cloning the pointer
+        # argument via `clone_preserve_strides`, which clones `storage_offset + numel` elements out
+        # of a buffer Inductor may have sized to the view alone -- an out-of-bounds read whose
+        # garbage survives into the loss on rows where the kernel returns early (ignore_index).
+        loss_1d_slice = torch.zeros(n_rows, dtype=loss_1d.dtype, device=device)  # chunk_size,
+        z_loss_1d_slice = torch.zeros(n_rows, dtype=z_loss_1d.dtype, device=device) if return_z_loss else None
+        token_accuracy_1d_slice = (
+            torch.zeros(n_rows, dtype=token_accuracy_1d.dtype, device=device) if return_token_accuracy else None
+        )
+        predicted_tokens_1d_slice = (
+            torch.full((n_rows,), -1, dtype=predicted_tokens_1d.dtype, device=device)
+            if return_predicted_tokens
+            else None
+        )
 
         # ensure _input and target are contiguous
         logits_chunk = logits_chunk.contiguous()

src/liger_kernel/ops/fused_linear_jsd.py

@@ -57,8 +57,11 @@ def fused_linear_jsd_forward(
             torch.zeros_like(student_weight, dtype=accum_dtype, device=device) if student_weight.requires_grad else None
         )
     grad_input = torch.zeros_like(student_input)
-    # we use fp32 for loss accumulator
-    loss_1d = torch.zeros((BT, V), dtype=torch.float32, device=device)
+    # we use fp32 for loss accumulator.
+    # Only the reduced loss is ever returned, so accumulate a scalar per chunk rather than
+    # materializing an unreduced BT x V buffer -- the latter costs as much as the full logits
+    # tensor that chunking exists to avoid (~4 GB at BT=8192, V=128000).
+    loss = torch.zeros((), dtype=torch.float32, device=device)
 
     if has_label:
         n_non_ignore = (shift_labels != ignore_index).sum().item()
@@ -80,8 +83,13 @@ def fused_linear_jsd_forward(
         teacher_logits_chunk = (teacher_input_chunk @ teacher_weight.t()).to(torch.float32)
         chunk_n_rows = student_logits_chunk.shape[0]
 
-        # unreduced loss
-        loss_1d_slice = loss_1d[start_idx:end_idx]  # chunk_size
+        # unreduced loss for this chunk
+        # NOTE: this must be a standalone contiguous buffer, *not* a view into a larger loss
+        # tensor. `torch.compile` functionalizes the kernel's mutation by cloning the pointer
+        # argument via `clone_preserve_strides`, which clones `storage_offset + numel` elements out
+        # of a buffer Inductor may have sized to the view alone -- an out-of-bounds read whose
+        # garbage survives into the loss on rows where the kernel returns early (ignore_index).
+        loss_chunk = torch.zeros((chunk_n_rows, V), dtype=torch.float32, device=device)
         # log-softmax with temperature
         student_logits_chunk = student_logits_chunk / temperature
         teacher_logits_chunk = teacher_logits_chunk / temperature
@@ -98,8 +106,8 @@ def fused_linear_jsd_forward(
             X_stride=student_prob_chunk.stride(-2),
             Y_ptr=teacher_prob_chunk,
             Y_stride=teacher_prob_chunk.stride(-2),
-            loss_ptr=loss_1d_slice,
-            loss_stride=loss_1d_slice.stride(-2),
+            loss_ptr=loss_chunk,
+            loss_stride=loss_chunk.stride(-2),
             dX_ptr=student_prob_chunk,
             dX_stride=student_prob_chunk.stride(-2),
             label_ptr=(
@@ -112,7 +120,7 @@ def fused_linear_jsd_forward(
             BLOCK_SIZE=BLOCK_SIZE,
             HAS_LABEL=has_label,
         )
-        loss_1d[start_idx:end_idx] = loss_1d_slice
+        loss = loss + loss_chunk.sum()
         # gradients of prob_chunk in place, shape: chunk_size x V
         # gradients of logits_chunk in place, shape: chunk_size x V
         student_logits_chunk = (
@@ -150,7 +158,6 @@ def fused_linear_jsd_forward(
                 else:
                     grad_weight.add_(torch.mm(grad_logits_t, student_input_chunk).float())
 
-    loss = torch.sum(loss_1d)
     grad_weight = (
         grad_weight.to(student_weight.dtype) if grad_weight is not None and accum_dtype is not None else grad_weight
     )

src/liger_kernel/ops/jsd.py

@@ -40,9 +40,14 @@ def _jsd_kernel(
     if HAS_LABEL:
         label = tl.load(label_ptr)
         if label == ignore_index:
+            # Write both outputs explicitly rather than relying on the caller having zeroed
+            # `loss_ptr`: under `torch.compile` the buffer handed to us is a functionalization
+            # clone, so caller-side initialization is not guaranteed to reach it.
             for i in range(0, n_cols, BLOCK_SIZE):
                 offsets = i + tl.arange(0, BLOCK_SIZE)
-                tl.store(dX_ptr + offsets, 0.0, mask=offsets < n_cols)
+                mask = offsets < n_cols
+                tl.store(loss_ptr + offsets, 0.0, mask=mask)
+                tl.store(dX_ptr + offsets, 0.0, mask=mask)
             return
 
     for i in range(0, n_cols, BLOCK_SIZE):

test/transformers/test_fused_linear_cross_entropy.py

@@ -1092,3 +1092,58 @@ def liger():
     assert_verbose_allclose(lig_gw, ref_gw, atol=atol, rtol=rtol)
     if bias:
         assert_verbose_allclose(lig_gb, ref_gb, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize(
+    "B, T, H, V",
+    [
+        (1, 256, 32, 4096),  # chunk_size == 2, so later chunks have large storage offsets
+        (2, 512, 64, 32000),
+    ],
+)
+@pytest.mark.parametrize("return_z_loss", [False, True])
+@torch.no_grad()
+def test_correctness_with_torch_compile(B, T, H, V, return_z_loss):
+    """`torch.compile` must not change the loss when some targets are `ignore_index`.
+
+    Regression test for silent numerical corruption (observed as `nan` losses): the per-chunk loss
+    buffers used to be `loss_1d[start_idx:end_idx]` views with non-zero storage offsets.
+    `torch.compile` functionalizes the kernel's mutation of them by cloning
+    `storage_offset + numel` elements out of a buffer Inductor may size to the view alone, i.e. an
+    out-of-bounds read. That garbage was normally overwritten by the kernel, but on an
+    `ignore_index` row `liger_cross_entropy_kernel` returns early without writing `loss_ptr`, so the
+    garbage survived into the reduced loss.
+
+    Forward-only: `fused_linear_cross_entropy_backward` branches on a tensor value, so the backward
+    pass cannot be captured with `fullgraph=True`.
+    """
+    ignore_index = -100
+    _input = torch.randn(B * T, H, device=device, dtype=torch.float32)
+    weight = torch.randn(V, H, device=device, dtype=torch.float32)
+
+    target = torch.randint(0, V, (B * T,), device=device, dtype=torch.long)
+    # at least one ignored target, but not all of them
+    target[torch.randperm(B * T)[: max(1, B * T // 4)]] = ignore_index
+
+    def call(_input, weight, target):
+        return LigerFusedLinearCrossEntropyFunction.apply(
+            _input,
+            weight,
+            target,
+            None,  # bias
+            None,  # ce_weight
+            ignore_index,
+            0.0,  # lse_square_scale
+            0.0,  # label_smoothing
+            "mean",  # reduction
+            None,  # softcap
+            return_z_loss,
+        )[0]
+
+    expected = call(_input, weight, target)
+
+    torch._dynamo.reset()
+    actual = torch.compile(call, fullgraph=True)(_input, weight, target)
+
+    # Identical kernels on identical inputs: the only permitted difference is reduction ordering.
+    assert_verbose_allclose(actual, expected, atol=1e-5, rtol=1e-5)

test/transformers/test_fused_linear_jsd.py

@@ -309,6 +309,66 @@ def test_correctness_functional(B, T, H, V, scalar, dtype, beta, ignore_index, t
     assert_verbose_allclose(_weight1.grad, _weight2.grad, atol=atol, rtol=rtol)
 
 
+@pytest.mark.parametrize(
+    "B, T, H, V",
+    [
+        (1, 4, 8, 2048),  # forces chunk_size == 1, i.e. one kernel launch per token
+        (2, 64, 128, 4096),
+    ],
+)
+@pytest.mark.parametrize(
+    "temperature, beta, ignore_index",
+    [
+        (1.0, 0.0, -100),
+        (1.0, 0.5, -100),
+        (2.0, 0.1, 42),
+        (1.0, 1.0, 2),
+    ],
+)
+@torch.no_grad()
+def test_correctness_with_torch_compile(B, T, H, V, beta, ignore_index, temperature):
+    """`torch.compile` must not change the loss when some labels are `ignore_index`.
+
+    Regression test for silent numerical corruption: the per-chunk loss buffer used to be a
+    `loss_1d[start_idx:end_idx]` view with a non-zero storage offset. `torch.compile` functionalizes
+    the kernel's mutation of it by cloning `storage_offset + numel` elements out of a buffer
+    Inductor may size to the view alone, i.e. an out-of-bounds read. That garbage was normally
+    overwritten by the kernel, but on an `ignore_index` row `_jsd_kernel` returns early without
+    writing `loss_ptr`, so the garbage survived into the reduced loss.
+
+    Forward-only: `fused_linear_jsd_backward` branches on a tensor value, so the backward pass
+    cannot be captured with `fullgraph=True`.
+    """
+    student_weight = torch.rand(V, H // 2, device=device, dtype=torch.float32)
+    teacher_weight = torch.rand(V, H, device=device, dtype=torch.float32)
+    student_input = torch.rand(B * T, H // 2, device=device, dtype=torch.float32)
+    teacher_input = torch.rand(B * T, H, device=device, dtype=torch.float32)
+
+    label = torch.randint(0, V, (B * T,), device=device, dtype=torch.long)
+    # at least one ignored label, but not all of them
+    label[torch.randperm(B * T)[: max(1, B * T // 4)]] = ignore_index
+
+    def call(student_input, teacher_input, label):
+        return liger_fused_linear_jsd(
+            student_input=student_input,
+            student_weight=student_weight,
+            teacher_input=teacher_input,
+            teacher_weight=teacher_weight,
+            shift_labels=label,
+            jsd_beta=beta,
+            ignore_index=ignore_index,
+            temperature=temperature,
+        )
+
+    expected = call(student_input, teacher_input, label)
+
+    torch._dynamo.reset()
+    actual = torch.compile(call, fullgraph=True)(student_input, teacher_input, label)
+
+    # Identical kernels on identical inputs: the only permitted difference is reduction ordering.
+    assert_verbose_allclose(actual, expected, atol=1e-5, rtol=1e-5)
+
+
 @pytest.mark.parametrize(
     "B, T, H, V",
     [