Lightning-AI / Lightning-AI/lightning-thunder
`AssertionError` for mutation on intermediates being reordered after `return`
@shino16 is already working on this.
Since Nov 29, 2025.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
### Bug
On `main`, 9e12768086aa6ef904927f6cc4760b8c8156fd32
```py
import torch, thunder, thunder.dynamo
def f():
torch.randn(4, device='cuda').sin_()
jf = thunder.dynamo.thunderfx(f)
# or jf = thunder.jit(f, fusion_type="dataflow")
jf()
```
```
File "/opt/pytorch/lightning-thunder/thunder/executors/nvfuserex_impl.py", line 928, in fusion_pass
fusedtrace = self.cse(fusedtrace)
^^^^^^^^^^^^^^^^^^^^
File "/opt/pytorch/lightning-thunder/thunder/executors/nvfuserex_impl.py", line 821, in cse
assert return_bsym.sym.id == prims.PrimIDs.RETURN
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError
```
`nvFuserExecutor` fuses the trace into
```py
# Constructed by Remove redundant casts (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation():
# tmp/main.py:4: torch.randn(4, device='cuda').sin_()
t6 = torch.randn((4,), device=torch.device("cuda:0"), dtype=torch.float32) # t6: "cuda:0 f32[4]"
# t6 = ltorch.randn((4,), generator=None, dtype=torch.float32, device=torch.device("cuda:0"), layout=torch.strided, requires_grad=False, pin_memory=False, out=None) # t6: "cuda:0 f32[4]"
# t6 = prims.randn((4,), device=devices.Device("cuda:0"), dtype=dtypes.float32) # t6: "cuda:0 f32[4]"
return {'output': (), 'flat_args': []}
(t7,) = update_aliases((t6,))
nvFusion0(t7)
# t1 = prims.sin(t7) # t1: "cuda:0 f32[4]"
# t2 = prims.copy_(t1, t7, grad_enabled=True) # t2: "cuda:0 f32[4]"
```
`sin_` here on intermediate `t7` is meaningless, but it survives through DCE because `prims.copy` has `DONT_DCE` tag. ~~(Afterwards it gets fused into `nvFusion0`, which gets DCE'd.)~~ `nvFuser0` will not be DCE'd because it has `copy_` as a subsymbol.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.