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
Conversation
`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>
This was referenced Aug 5, 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>
Contributor
Author
|
Any possibility of landing this? It would help overcome the bug in pytorch (pytorch/pytorch#192227) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
LigerFusedLinearJSDandLigerFusedLinearCrossEntropyreturn a silently wrong loss undertorch.compilewhenever the labels contain at least oneignore_indexposition. There is no exception, warning, or graph-break error — the compiled graph runs to completion and returns a numerically incorrect value. ForLigerFusedLinearCrossEntropythe loss becomesnanat common shapes.Reproduces with both
torch.compile(fn, fullgraph=True)and plaintorch.compile(fn). An all-valid mask (zero ignored positions) does not reproduce it.Commits
This PR has two independent, separately revertable commits:
50f4a69torch.compilecorrectness bug described below4b77c6bBT x Vfp32 loss buffer infused_linear_jsdThey 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 rawtriton.jitkernel.torch.compilefunctionalizes that mutation withclone_preserve_strides, which must clonestorage_offset + numelelements so the kernel's pointer arithmetic still lands correctly: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_stridedas an unmasked copy — reading past the end. FromTORCH_LOGS=output_codeon the repro above (slice_15isf32[1, 2048], but 4096 elements are cloned out of it):The generated buffer sizes match
needed_sizeexactly 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_indexrow both kernels return early without writingloss_ptr, so the garbage survives intotorch.sum(loss_1d).Both preconditions are required, which is exactly why an all-valid mask is unaffected — every
loss_ptrgets fully written, so the garbage is always overwritten.Confirming it is foreign memory rather than a logic error:
compiled - eagerwas preciselyrow0_loss / 3, i.e. row 0's contribution counted twice, and settingtorch._inductor.config.allow_buffer_reuse = Falsechanged the wrong value to104.598.compute-sanitizerreports nothing, because the over-read stays inside the caching allocator's segment.LigerFusedLinearCrossEntropyis affected too, and worseSame pattern at
fused_linear_cross_entropy.py, whereliger_cross_entropy_kernelskips bothloss_ptrandz_loss_ptron the ignore path. Becausechunk_sizeshrinks asHshrinks, later chunks get much larger storage offsets and the over-read is far bigger:BT,H,V)z_loss7 of 20 swept configurations were wrong before this change; 0 of 20 after.
Memory: dropping the
BT x Vloss buffer (commit4b77c6b)Noticed while fixing the above.
fused_linear_jsd_forwardallocatedloss_1das an unreducedBT x Vfp32 tensor, wrote each chunk's slice into it, then reduced the whole thing with a singletorch.sum. Nothing else ever read it — so the intermediate cost as much as the full logits tensorthat 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:
BT=1024, H=512, V=32000BT=4096, H=1024, V=32000BT=2048, H=512, V=128000This 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:BT=1024, V=32000, beta=0.5BT=4096, V=32000, beta=0.0BT=2048, V=128000, beta=0.5Gradients are unchanged — the gradient path never touched this buffer. All existing JSD tests pass
against the PyTorch reference at their current tolerances (
133 passedacrosstest_jsd.pyandtest_fused_linear_jsd.py), so no tolerance changes were needed.Side benefit: eager and
torch.compilenow agree bit-for-bit on the repro above(
5.865699291229248both), since both reduce in the same order.If the reduction-order change is unwelcome in a correctness PR, drop
4b77c6b—50f4a69standsalone.
The fix
Two layers, each independently sufficient (verified by reverting one at a time):
needed_size == numeland the functionalization clone is exact. This also makes the pre-existingloss_1d[start_idx:end_idx] = loss_1d_slicewrite-back meaningful rather than a self-assignment._tv_distance_kernelalready does exactly this, so this brings JSD and CE in line with existing convention.Eager numerics are unchanged — the same
BT x Vbuffer is reduced in the same order.Not changed
element_mul_kernelwhole 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 (maxdiff0.0); fed offset views the same kernel diverges by9.286, confirming the calling convention is the discriminator. End-to-end gradients also match eager withignore_indexpresent, for bothgrad_output == 1.0(whereelement_mul_kernelis skipped) andgrad_output == 2.5(where it runs)._ascend/ops/fused_linear_cross_entropy.py(offset views) and its generic CE kernel's ignore path, which writeslse/token_accuracy/predicted_tokensbut notloss_ptr/z_loss_ptr. Its JSD kernel already zeroesloss_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.as_stridedbeyond 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
if torch.ne(grad_output, ...)/torch.equal), sofullgraph=Truecannot capture the backward pass. Relatedly,torch._dynamo.compiled_autogradfails outright on both withaten::_local_scalar_dense() ... Unable to cast 1.0 to Tensor.Testing Done
torch 2.13.0+cu130,triton 3.7.1, Python 3.12.3make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergenceTargeted regression tests
Both new tests were verified to fail without the fix and pass with it:
make checkstylemake testRun on both commits of this PR and, for comparison, on
main@b1962fa:main@b1962fa50f4a69(fix)4b77c6b(fix + perf)The 2 failures are pre-existing on
main— same two parametrizations, NaN inlogits_triton.grad. Unrelated to this change (ops/grpo_loss.pyimports onlytorch/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, andops/grpo_loss.pyhas several barereturns in its kernels. Possibly worth the same audit, in a separate issue.make test-convergenceRun on
50f4a69. Not re-run on4b77c6b: that commit touches onlyops/fused_linear_jsd.py, andtest/convergence/contains zero references to JSD, so theconvergence suite cannot reach it. The test files that do reach it were re-run (see above).
None are caused by this change. Each was re-run on
main@b1962fa:mainfp32mini_llavafp32mini_nemotronbf16_with_logitsmini_llama4bf16_with_logitsmini_qwen3_moebf16mini_qwen3_moebf16mini_qwen3_5_moemain, 1 in 4 on this PRThe 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:
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 store0.0where the caller had already stored0.0. Only thetorch.compilefunctionalization path differs.Re-run on
4b77c6b(the perf commit)Full
make testwas re-run on this commit (see the table above: identical counts to50f4a69,same two pre-existing failures). That commit touches only
ops/fused_linear_jsd.py, and every testfile that reaches it was also run individually: