Lightning-AI / Lightning-AI/lightning-thunder

`AssertionError` for mutation on intermediates being reordered after `return`

Open
#2,776 1 comment 0 reactions 1 assignee View on GitHub

@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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.