Lightning-AI / Lightning-AI/lightning-thunder
thunder.jit does not allow changing the model structure (None->Buffer) in forward and leads to incorrect output if doing so
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
This sort of code occurs in `BAAI/Bunny-v1_1-4B`, see https://huggingface.co/BAAI/Bunny-v1_1-4B/blob/2faf1780f240d4fc560d61142160a8562bfb6477/modeling_bunny_phi3.py#L1099-L1126
Found while investigating #1950
```python
import torch
import thunder
class Model(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.fc = torch.nn.Linear(2, 2)
self.register_buffer("buf", None, False)
def forward(self, x):
if self.buf is None:
self.buf = torch.randn(1, 2, device=x.device)
return self.fc(x) + self.buf
with torch.device("cuda"):
org_m = Model()
# m = org_m
m = thunder.jit(org_m)
x = torch.randn(2, 2)
torch.testing.assert_close(m(x), org_m(x))
# With eager, tensor([[1., 1.]], device='cuda:0')
# With thunder, None
print(org_m.buf)
m(x)
thunder.last_traces(m)[-1].save_trace("trace.py")
torch.testing.assert_close(m(x), org_m(x)) # Errors
```
Generated Trace (see that torch.randn will be called on every invocation)
```python
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
import torch.nn.functional
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(x, t_fc_bias, t_fc_weight):
# x: "cuda:0 f32[2, 2]"
# t_fc_bias: "cuda:0 f32[2]"
# t_fc_weight: "cuda:0 f32[2, 2]"
t36 = torch.randn((1, 2), device=torch.device("cuda:0"), dtype=torch.float32) # t36: "cuda:0 f32[1, 2]"
# t36 = ltorch.randn((1, 2), generator=None, dtype=torch.float32, device=torch.device("cuda:0"), layout=torch.strided, requires_grad=False, pin_memory=False, out=None) # t36: "cuda:0 f32[1, 2]"
# t36 = prims.randn((1, 2), device=devices.Device("cuda:0"), dtype=dtypes.float32) # t36: "cuda:0 f32[1, 2]"
# /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:125: return F.linear(input, self.weight, self.bias)
t37 = torch.nn.functional.linear(x, t_fc_weight, t_fc_bias) # t37: "cuda:0 f32[2, 2]"
# t37 = ltorch.linear(x, t_fc_weight, t_fc_bias) # t37: "cuda:0 f32[2, 2]"
# t37 = prims.linear(x, t_fc_weight, t_fc_bias) # t37: "cuda:0 f32[2, 2]"
[t30] = nvFusion0(t36, t37)
# t34 = prims.broadcast_in_dim(t36, (2, 2), (0, 1)) # t34: "cuda:0 f32[2, 2]"
# t30 = prims.add(t37, t34) # t30: "cuda:0 f32[2, 2]"
del t37
return {'output': (t36, t30), 'flat_args': [x, t_fc_bias, t_fc_weight], 'flat_output': (t36, t30)}, ((x,), ())
```
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.
Research direction
Start with the minimal Python reproduction in the issue and compare eager execution with thunder.jit, then inspect the saved trace from thunder.last_traces(m)[-1].save_trace("trace.py"). Done means a buffer initialized from None during forward is retained across calls and the eager and compiled outputs match without regenerating it each invocation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100