Skip to content

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

Open
KyleMylonakisProtopia wants to merge 5 commits into
linkedin:mainfrom
KyleMylonakisProtopia:fix/torch-compile-ignore-index-loss-corruption
Open

fix: torch.compile loss corruption at ignore_index rows (+ perf: drop BT x V loss buffer in FLJSD)#1338
KyleMylonakisProtopia wants to merge 5 commits into
linkedin:mainfrom
KyleMylonakisProtopia:fix/torch-compile-ignore-index-loss-corruption

Conversation

@KyleMylonakisProtopia

@KyleMylonakisProtopia KyleMylonakisProtopia commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Summary

LigerFusedLinearJSD and LigerFusedLinearCrossEntropy return a silently wrong loss under torch.compile whenever the labels contain at least one ignore_index position. There is no exception, warning, or graph-break error — the compiled graph runs to completion and returns a numerically incorrect value. For LigerFusedLinearCrossEntropy the loss becomes nan at common shapes.

Reproduces with both torch.compile(fn, fullgraph=True) and plain torch.compile(fn). An all-valid mask (zero ignored positions) does not reproduce it.

import torch
from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD

torch.manual_seed(0)
student, teacher = torch.randn(4, 8, device="cuda"), torch.randn(4, 8, device="cuda")
weight = torch.randn(2048, 8, device="cuda")
fused_jsd = LigerFusedLinearJSD(jsd_beta=0.0, ignore_index=-100, temperature=1.0)
shift_labels = torch.tensor([0, -100, 0, 0], device="cuda")   # one ignored position

call = lambda s, t, l: fused_jsd(s, weight, t, weight, l)
print(call(student, teacher, shift_labels).item())                                  # 5.8657
print(torch.compile(call, fullgraph=True)(student, teacher, shift_labels).item())    # 9.1759  <-- wrong

Commits

This PR has two independent, separately revertable commits:

commit what
50f4a69 fix — the torch.compile correctness bug described below
4b77c6b perf — stop materializing an unreduced BT x V fp32 loss buffer in fused_linear_jsd

They are independent: either can land alone. Happy to split the second one into its own PR if preferred.

Details

Root cause

Both ops pass loss_1d[start_idx:end_idx] — a view with a non-zero storage offset — as a mutated pointer to a raw triton.jit kernel.

torch.compile functionalizes that mutation with clone_preserve_strides, which must clone storage_offset + numel elements so the kernel's pointer arithmetic still lands correctly:

needed_size = compute_required_storage_length(x.size(), x.stride(), x.storage_offset())
buffer = torch.as_strided(x, (needed_size,), (1,), 0).clone()
return torch.as_strided(buffer, x.size(), x.stride(), x.storage_offset())

That is correct in eager, where the real base storage is large enough. But Inductor may realize the view as a buffer sized to the view alone, and then lowers the as_strided as an unmasked copy — reading past the end. From TORCH_LOGS=output_code on the repro above (slice_15 is f32[1, 2048], but 4096 elements are cloned out of it):

%as_strided_default_4 : "f32[4096][1]" = aten.as_strided.default(%slice_15, [4096], [1], 0)
%clone_default_2      : "f32[4096][1]" = aten.clone.default(%as_strided_default_4)
def triton_poi_fused_5(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 4096                                  # in_ptr0 (buf24) holds only 2048 elements
    xmask = tl.full([XBLOCK], True, tl.int1)[:]    # unmasked -> 2048 elements read out of bounds

The generated buffer sizes match needed_size exactly for each chunk: 4096 / 6144 / 8192 for rows 1 / 2 / 3.

That over-read is normally harmless, because the kernel overwrites the whole slice. But on an ignore_index row both kernels return early without writing loss_ptr, so the garbage survives into torch.sum(loss_1d).

Both preconditions are required, which is exactly why an all-valid mask is unaffected — every loss_ptr gets fully written, so the garbage is always overwritten.

Confirming it is foreign memory rather than a logic error: compiled - eager was precisely row0_loss / 3, i.e. row 0's contribution counted twice, and setting torch._inductor.config.allow_buffer_reuse = False changed the wrong value to 104.598. compute-sanitizer reports nothing, because the over-read stays inside the caching allocator's segment.

LigerFusedLinearCrossEntropy is affected too, and worse

Same pattern at fused_linear_cross_entropy.py, where liger_cross_entropy_kernel skips both loss_ptr and z_loss_ptr on the ignore path. Because chunk_size shrinks as H shrinks, later chunks get much larger storage offsets and the over-read is far bigger:

shape (BT, H, V) eager compiled (before) compiled (after)
256, 32, 4096 21.087738 nan 21.087738
1024, 64, 32000 32.996323 nan 32.996323
1024, 64, 32000, z_loss 32.996323 36.991989 32.996323

7 of 20 swept configurations were wrong before this change; 0 of 20 after.

Memory: dropping the BT x V loss buffer (commit 4b77c6b)

Noticed while fixing the above. fused_linear_jsd_forward allocated loss_1d as an unreduced
BT x V fp32 tensor, wrote each chunk's slice into it, then reduced the whole thing with a single
torch.sum. Nothing else ever read it — so the intermediate cost as much as the full logits tensor
that chunking exists to avoid: ~4 GB at BT=8192, V=128000.

It now accumulates a scalar per chunk. Measured peak-allocation reduction, matching the removed
buffer's footprint exactly:

shape before after delta
BT=1024, H=512, V=32000 672.2 MiB 546.2 MiB -126 MiB
BT=4096, H=1024, V=32000 2328.5 MiB 1828.5 MiB -500 MiB
BT=2048, H=512, V=128000 4527.5 MiB 3527.5 MiB -1000 MiB

This changes the reduction order — per-chunk partial sums accumulated sequentially, instead of
one pass over BT x V. Measured impact on the loss, at fp32 epsilon scale:

shape before after rel. diff
BT=1024, V=32000, beta=0.5 0.2009693086 0.2009692937 7.4e-8
BT=4096, V=32000, beta=0.0 1.0008158684 1.0008158684 0 (bit-identical)
BT=2048, V=128000, beta=0.5 0.2013413161 0.2013413757 3.0e-7

Gradients are unchanged — the gradient path never touched this buffer. All existing JSD tests pass
against the PyTorch reference at their current tolerances (133 passed across test_jsd.py and
test_fused_linear_jsd.py), so no tolerance changes were needed.

Side benefit: eager and torch.compile now agree bit-for-bit on the repro above
(5.865699291229248 both), since both reduce in the same order.

If the reduction-order change is unwelcome in a correctness PR, drop 4b77c6b50f4a69 stands
alone.

The fix

Two layers, each independently sufficient (verified by reverting one at a time):

  1. Allocate standalone contiguous per-chunk output buffers instead of views, so needed_size == numel and the functionalization clone is exact. This also makes the pre-existing loss_1d[start_idx:end_idx] = loss_1d_slice write-back meaningful rather than a self-assignment.
  2. Write zeros to the loss pointers on the ignore path, so neither kernel relies on the caller having zeroed a buffer it does not own. _tv_distance_kernel already does exactly this, so this brings JSD and CE in line with existing convention.

Eager numerics are unchanged — the same BT x V buffer is reduced in the same order.

Not changed

  • The backward pass is unaffected. It only ever hands element_mul_kernel whole tensors at storage offset 0, so the clone is exact, and that kernel has no early-return path. Verified in isolation under Inductor lowering: fed whole tensors it matches eager exactly (maxdiff 0.0); fed offset views the same kernel diverges by 9.286, confirming the calling convention is the discriminator. End-to-end gradients also match eager with ignore_index present, for both grad_output == 1.0 (where element_mul_kernel is skipped) and grad_output == 2.5 (where it runs).
  • The Ascend backend carries the same pattern in _ascend/ops/fused_linear_cross_entropy.py (offset views) and its generic CE kernel's ignore path, which writes lse / token_accuracy / predicted_tokens but not loss_ptr / z_loss_ptr. Its JSD kernel already zeroes loss_row_ptr. I have no NPU to test on, so I left it alone rather than ship an unverified change to a hardware backend — flagging it here for someone with access.
  • Arguably also an upstream Inductor bug, since it lowers as_strided beyond a realized buffer's extent as an unmasked read. Worth reporting separately; this change removes Liger's dependence on that behavior either way.

Adjacent issues found, worth separate tickets

  • Both backward functions branch on a tensor value (if torch.ne(grad_output, ...) / torch.equal), so fullgraph=True cannot capture the backward pass. Relatedly, torch._dynamo.compiled_autograd fails outright on both with aten::_local_scalar_dense() ... Unable to cast 1.0 to Tensor.

Testing Done

  • Hardware Type: NVIDIA GB10 (compute capability 12.1, 48 SMs)torch 2.13.0+cu130, triton 3.7.1, Python 3.12.3
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

Targeted regression tests

Both new tests were verified to fail without the fix and pass with it:

# without the fix
FAILED test_fused_linear_jsd.py::test_correctness_with_torch_compile[1.0-0.0--100-2-64-128-4096]
FAILED test_fused_linear_jsd.py::test_correctness_with_torch_compile[1.0-0.5--100-2-64-128-4096]
FAILED test_fused_linear_jsd.py::test_correctness_with_torch_compile[2.0-0.1-42-2-64-128-4096]
FAILED test_fused_linear_jsd.py::test_correctness_with_torch_compile[1.0-1.0-2-2-64-128-4096]
4 failed, 4 passed
FAILED test_fused_linear_cross_entropy.py::test_correctness_with_torch_compile[False-1-256-32-4096]
FAILED test_fused_linear_cross_entropy.py::test_correctness_with_torch_compile[False-2-512-64-32000]
FAILED test_fused_linear_cross_entropy.py::test_correctness_with_torch_compile[True-1-256-32-4096]
FAILED test_fused_linear_cross_entropy.py::test_correctness_with_torch_compile[True-2-512-64-32000]
4 failed

# with the fix
8 passed   (test_fused_linear_jsd.py -k test_correctness_with_torch_compile)
4 passed   (test_fused_linear_cross_entropy.py -k test_correctness_with_torch_compile)

# affected kernels' full suites
463 passed  (test_cross_entropy.py test_fused_linear_cross_entropy.py test_jsd.py test_fused_linear_jsd.py)

make checkstyle

All checks passed!
317 files already formatted

make test

Run on both commits of this PR and, for comparison, on main @ b1962fa:

main @ b1962fa 50f4a69 (fix) 4b77c6b (fix + perf)
passed 3923 3935 (+12 = the new tests) 3935
failed 2 2 (same node IDs) 2 (same node IDs)
skipped 923 923 923
xfailed 14 14 14
# this PR @ 50f4a69
= 2 failed, 3935 passed, 923 skipped, 14 xfailed, 80 warnings in 1751.69s (0:29:11) =
FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-False-luspo-sequence-1.5]
FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-True-luspo-sequence-1.5]

# this PR @ 4b77c6b
= 2 failed, 3935 passed, 923 skipped, 14 xfailed, 80 warnings in 826.06s (0:13:46) =
FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-False-luspo-sequence-1.5]
FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-True-luspo-sequence-1.5]

# main @ b1962fa (baseline)
= 2 failed, 3923 passed, 923 skipped, 14 xfailed, 79 warnings in 794.06s (0:13:14) =
FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-False-luspo-sequence-1.5]
FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-True-luspo-sequence-1.5]

The 2 failures are pre-existing on main — same two parametrizations, NaN in logits_triton.grad. Unrelated to this change (ops/grpo_loss.py imports only torch/triton). They pass in isolation and only fail in full-suite context.

Side note: those NaNs sit in a masked row of completion_mask, which is the same shape of defect this PR fixes, and ops/grpo_loss.py has several bare returns in its kernels. Possibly worth the same audit, in a separate issue.

make test-convergence

Run on 50f4a69. Not re-run on 4b77c6b: that commit touches only
ops/fused_linear_jsd.py, and test/convergence/ contains zero references to JSD, so the
convergence suite cannot reach it. The test files that do reach it were re-run (see above).

= 6 failed, 153 passed, 9 skipped, 3 xfailed, 1 xpassed, 307 warnings in 2972.03s (0:49:32) =
FAILED fp32/test_mini_models.py::test_mini_model[mini_llava-...]
FAILED fp32/test_mini_models.py::test_mini_model[mini_nemotron-...]
FAILED bf16/test_mini_models.py::test_mini_model[mini_qwen3_moe-...]
FAILED bf16/test_mini_models.py::test_mini_model[mini_qwen3_5_moe-...]
FAILED bf16/test_mini_models_with_logits.py::test_mini_model[mini_llama4-...]
FAILED bf16/test_mini_models_with_logits.py::test_mini_model[mini_qwen3_moe-...]

None are caused by this change. Each was re-run on main @ b1962fa:

test on main
fp32 mini_llava fails
fp32 mini_nemotron fails
bf16_with_logits mini_llama4 fails
bf16_with_logits mini_qwen3_moe fails
bf16 mini_qwen3_moe fails (2/2)
bf16 mini_qwen3_5_moe flaky — 1 failure in 4 runs on main, 1 in 4 on this PR

The MoE convergence tests are nondeterministic run-to-run. Re-running the same commit produces different values, with the reference side stable and only the Liger side moving:

main @ b1962fa run 1:  Mismatch at (0,31): tensor1 = 3.5302677, tensor2 = 4.730206
main @ b1962fa run 2:  Mismatch at (0,31): tensor1 = 3.5302677, tensor2 = 4.701219

That localizes the nondeterminism to the Liger MoE path (likely atomics / reduction order in ops/fused_moe_kernels.py) rather than to bf16 training noise. Also worth a separate issue.

This change cannot affect convergence in the first place: convergence tests run eager, and in eager the change is bit-identical — a standalone torch.zeros(n_rows) buffer holds the same values as the zeroed slice it replaces, and the kernels now store 0.0 where the caller had already stored 0.0. Only the torch.compile functionalization path differs.

Re-run on 4b77c6b (the perf commit)

Full make test was re-run on this commit (see the table above: identical counts to 50f4a69,
same two pre-existing failures). That commit touches only ops/fused_linear_jsd.py, and every test
file that reaches it was also run individually:

133 passed   test/transformers/test_jsd.py test/transformers/test_fused_linear_jsd.py
64 passed, 1 skipped   test/chunked_loss/test_jsd_loss.py test/transformers/test_cutile_backend.py
330 passed   test/transformers/test_cross_entropy.py test/transformers/test_fused_linear_cross_entropy.py  (no collateral)
All checks passed! / 317 files already formatted   (make checkstyle)

KyleMylonakisProtopia and others added 4 commits July 24, 2026 13:59
`fused_linear_jsd_forward` and `fused_linear_cross_entropy_forward` passed
`loss_1d[start_idx:end_idx]` -- a view with a non-zero storage offset -- as a
mutated pointer to a raw `triton.jit` kernel. `torch.compile` functionalizes
such a mutation with `clone_preserve_strides`, which emits
`as_strided(view, [storage_offset + numel], [1], 0)`. Inductor may realize the
view as a buffer sized to the view alone, so that clone reads past its end,
unmasked.

The over-read garbage was normally harmless because the kernel overwrote it.
But on an `ignore_index` row both kernels return early *without* writing
`loss_ptr`, so the garbage survived into the reduced loss. The result was a
silently wrong loss under `torch.compile` -- no exception and no graph-break
warning -- whenever `shift_labels`/`target` contained at least one
`ignore_index` position, and `nan` for `LigerFusedLinearCrossEntropy` at common
shapes.

Fix at both layers, either of which is independently sufficient:

- Allocate standalone contiguous per-chunk output buffers instead of views, so
  the functionalization clone is exact.
- Have `_jsd_kernel` and `liger_cross_entropy_kernel` write zeros to their loss
  pointers on the ignore path, so neither relies on the caller having zeroed a
  buffer it does not own. `_tv_distance_kernel` already does this.

Eager numerics are unchanged. The backward pass is unaffected: it only ever
hands `element_mul_kernel` whole tensors at storage offset 0, so the clone is
exact there, and that kernel has no early-return path.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Kyle Mylonakis <kyle@protopia.ai>
`fused_linear_jsd_forward` allocated `loss_1d` as an unreduced `BT x V` fp32
tensor, wrote each chunk's slice into it, and then reduced the whole thing with
a single `torch.sum`. Nothing else read it, so the intermediate buffer cost as
much as the full logits tensor that chunking exists to avoid -- roughly 4 GB at
`BT=8192, V=128000`, and 1 GB at `BT=2048, V=128000`.

Accumulate a scalar per chunk instead. Measured peak-allocation reduction,
exactly matching the removed buffer's footprint:

  BT=1024  H=512   V=32000    672.2 MiB -> 546.2 MiB   (-126)
  BT=4096  H=1024  V=32000   2328.5 MiB -> 1828.5 MiB  (-500)
  BT=2048  H=512   V=128000  4527.5 MiB -> 3527.5 MiB  (-1000)

This changes the reduction order: per-chunk partial sums are accumulated
sequentially rather than reduced in one pass over `BT x V`. Loss values move at
fp32 epsilon scale (largest observed relative difference 3.0e-7, several shapes
bit-identical) and gradients are unchanged, since the gradient path never
touched this buffer.

It also makes eager and `torch.compile` agree bit-for-bit on the repro from the
previous commit, as both now reduce in the same order.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Kyle Mylonakis <kyle@protopia.ai>
@KyleMylonakisProtopia KyleMylonakisProtopia changed the title fix: prevent torch.compile from corrupting loss at ignore_index rows fix: torch.compile loss corruption at ignore_index rows (+ perf: drop BT x V loss buffer in FLJSD) Aug 4, 2026
KyleMylonakisProtopia added a commit to KyleMylonakisProtopia/pytorch that referenced this pull request Aug 5, 2026
… views

Passing a view with a non-zero storage_offset as a mutated pointer argument to a
raw triton.jit kernel silently produced wrong results under torch.compile:
elements the kernel chose not to write came back holding data read from past the
end of a buffer instead of their original values. No exception, no warning, no
graph break. Found in the wild in Liger-Kernel, where it corrupted training
losses and produced NaNs (see linkedin/Liger-Kernel#1338).

The root cause is a contract mismatch between two components. Functionalizing a
mutated Triton arg clones it with clone_preserve_strides(), which allocates
storage_offset + span elements and copies them out of the tensor's storage via
as_strided(x, (needed_size,), (1,), 0). That is correct in eager, where the whole
storage is addressable. Inductor, though, materializes an intermediate view as a
buffer sized to the view alone, so the leading storage_offset elements lie
outside any allocation: aten.as_strided offsets are storage-relative, but
Inductor's as_strided lowering interprets them relative to whichever buffer got
realized, and built a ReinterpretView describing memory past the end with no
extent check. For a (4, 2048) base that emitted an unmasked copy kernel reading
4096 elements out of a 2048-element buffer, and the user's kernel then received a
pointer into precisely the garbage half.

The fix has two halves. The clone no longer asks for storage the tensor does not
own: for a dense tensor at a non-zero offset, a plain clone() reproduces its
strides and lands at storage_offset 0, which is all a triton.jit kernel needs,
since it only ever sees a data pointer plus the strides passed to it. A tensor
with gaps between its elements needs those gaps materialized too, so it is cloned
through its base, which clone_preserve_strides now accepts as an opt-in keyword.
The base is captured at the functionalization boundary and threaded through the
functional HOP, because reading _base inside the functional impl yields an
untracked tensor that the tracer lifts to a constant. Separately, Inductor's
as_strided lowering now refuses to reinterpret past a realized buffer rather than
silently doing so, so any remaining or future producer of this pattern fails
loudly instead of corrupting data.

Suggested review order: torch/_prims_common (the opt-in base), then
triton_kernel_wrap (which clone is chosen, and how the base is threaded), then
lowering.py (the backstop), then post_grad.py (keeps the new HOP kwarg off the
mutation node), then the tests.

Two alternatives were considered and rejected. Changing clone_preserve_strides'
default semantics to clone compactly would break prims.as_strided_scatter, whose
documented contract is "clone then as_strided(...).copy_()" with a
storage-relative offset, and would make the runtime callers (cudagraphs, compiled
autograd) copy whole bases; hence an opt-in keyword instead. Building the compact
clone as empty_strided(size, stride).copy_(val) reads more simply but is wrong:
empty_strided is a factory with no tensor input, so two clones of the same tensor
are identical calls that get CSE'd into a single buffer, and a kernel with two
mutated args then writes both of its outputs to the same memory.

The clones themselves cannot simply be dropped. They are what makes the
functional HOP functional, and eliding them is the reinplacing pass's job; that
pass handles a view of a graph input but bails for an intra-graph base, where
functionalization emits slice_scatter rather than copy_ and there is no copy-back
node to key on.

Test Plan:

New tests, each confirmed to fail with the corresponding half of the fix reverted
and to pass with it restored:

```
python test/inductor/test_triton_kernels.py -k mutated_offset_view
python test/inductor/test_torchinductor.py CpuTests.test_as_strided_past_input_extent_cpu GPUTests.test_as_strided_past_input_extent_cuda CpuTests.test_as_strided_past_realized_buffer_raises_cpu GPUTests.test_as_strided_past_realized_buffer_raises_cuda
```

test_as_strided_past_input_extent pins the InputBuffer exemption in the new
check: a graph input aliasing a larger storage may legitimately be as_strided
past its own extent, because Inductor rebases the offset onto the incoming
pointer, and the check must not reject that. The gap=2 parametrization covers a
non-dense view, whose required extent exceeds its numel even at offset 0.
test_triton_kernel_mutated_offset_view_large_tensor uses a base of 2147485696
elements to confirm Inductor promotes its generated kernels to 64-bit indexing
past int32 max; it is gated with largeTensorTest("6GB").

Regression suites:

```
python test/inductor/test_triton_kernels.py
python test/inductor/test_auto_functionalize.py
python test/inductor/test_torchinductor.py -k as_strided
python test/inductor/test_torchinductor_dynamic_shapes.py -k as_strided
lintrunner -a
```

test_triton_kernels.py runs 428 tests with 2 failures, both
test_tma_capture_and_functionalize_*_tma_version_new, which fail identically with
this change reverted. Everything else passes.

Also checked by hand against a deliberately poisoned caching allocator, so that
corruption would be identifiable rather than incidental: static shapes,
dynamic=True, inference_mode, no_grad, a non-dense offset view, and a mutable
custom op taking an offset view all return correct results with no freed-memory
values present.

This change was authored with the assistance of an AI coding assistant (Claude
Code). Claude Opus 5 (1M context) <noreply@anthropic.com>

Signed-off-by: Kyle Mylonakis <kyle@protopia.ai>
@KyleMylonakisProtopia

Copy link
Copy Markdown
Contributor Author

Any possibility of landing this? It would help overcome the bug in pytorch (pytorch/pytorch#192227)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant