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

Open
#1,951 1 comment 0 reactions 0 assignees View on GitHub

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

  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.

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.