$ the-wire · showcase
Inductor fixes stale captured-tensor mutations in while_loop
By RepoJournal · Filed · About PyTorch
15 people shipped this
Jason Ansel
@jansel
1 cited
Aditya Sharma
@adityasharmaaaaa
1 cited
nithinsubbiah
@nithinsubbiah
1 cited
ivy-zhou
@ivy-zhou
1 cited
acisseJZhong
@acisseJZhong
1 cited
Nikita Shulga
@malfet
srhust
@srhust
Anthony Shoumikhin
@shoumikhin
Sebastian-Larsson
@Sebastian-Larsson
YufengShi-dudu
@YufengShi-dudu
SaoirseARM
@SaoirseARM
pytorchbot
@pytorchbot
vmoens
@vmoens
ゆり
@yurekami
wwwjn
@wwwjn
PyTorch's Inductor now correctly replays cond_fn mutations of captured tensors in while_loop, fixing silent wrong results in later compiled calls.
Inductor fixed a correctness bug in `torch.while_loop` [1]. Under inference mode, when `cond_fn` mutated a captured tensor that was not a carried input and that tensor was mutated before the loop, the first compiled call returned correct values but left the external tensor stale, and later calls silently produced wrong results. The fix replays these mutations, so external storage stays consistent across calls.
A test fix [2] corrects a grid-size error in `test_graph_partition_user_defined_triton_kernel_reuse` that caused compiled comparisons to check uninitialized memory. The kernel `add_kernel[(4,)](..., n_elements=128, BLOCK_SIZE=16)` launches only 64 Triton programs, leaving elements 64 to 127 unwritten; the test now launches enough programs to cover the full output.
A document fix [3] clarifies the actual requirement for `ROWS_GUARANTEED_SAFE` in the attention kernel: each row needs an unmasked key in the first block scheduled for it, not merely somewhere in the row, because the kernel runs a row maximum starting from `-inf`.
Elsewhere, TorchTitan now routes saves through the `torch_checkpointing` backend [4], which previously built its backend but raised `NotImplementedError` on every operation. Also, TorchTitan removed the in-tree `torch.experiments.forge` module [5]; users importing it from `main` should switch to the standalone TorchForge project.
Action items
- → Upgrade PyTorch to a build including commit fc660a8 to fix silent wrong results from while_loop with captured tensor mutations. pytorch/pytorch [immediate]
- → Review any custom Triton kernels that rely on n_elements to size the grid; verify the grid covers the full output range. pytorch/pytorch [plan]
- → If you import torchtitan.experiments.forge from TorchTitan main, migrate to the standalone TorchForge project. pytorch/torchtitan [plan]
References
- [1] [inductor] Replay cond_fn mutations of captured tensors in while_loop (#195393) ↗ pytorch/pytorch
- [2] [inductor] Fix incorrect grid size in test_graph_partition_user_defined_triton_kernel_reuse (#195546) ↗ pytorch/pytorch
- [3] [Inductor] Document the actual ROWS_GUARANTEED_SAFE requirement (#194154) ↗ pytorch/pytorch
- [4] Route saves through the torch_checkpointing backend ↗ pytorch/torchtitan
- [5] Deprecate Forge under experimental ↗ pytorch/torchtitan
