Lightning-AI / Lightning-AI/lightning-thunder
Updating Module's buffer leads to crash in thunder.jit
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
```python
import torch
import thunder
class TestModel(torch.nn.Module):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.register_buffer("buffer", torch.tensor([1,]))
def forward(self, x):
self.buffer = self.buffer + 1
return x + self.buffer
m = TestModel()
x = torch.randn(3, )
print(m(x))
jm = thunder.jit(TestModel())
print(jm(x))
```
Errors with:
```python
Traceback (most recent call last):
File "lightning-thunder/scratchpad/test.py", line 73, in
print(jm(x))
File "git/pytorch/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "git/pytorch/torch/nn/modules/module.py", line 1541, in _call_impl
return forward_call(*args, **kwargs)
File "lightning-thunder/thunder/core/module.py", line 61, in forward
res = self._forward_fn(*args, **kwargs)
File "lightning-thunder/thunder/__init__.py", line 675, in fn_
cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
File "lightning-thunder/thunder/__init__.py", line 224, in cache_info_wrapper
res = fn(*args, **kwargs)
File "lightning-thunder/thunder/__init__.py", line 503, in get_computation_and_inputs
jit_results: TraceResults = interpreter(
File "lightning-thunder/thunder/__init__.py", line 212, in _general_frontend
return thunder_general_jit(fn, args, kwargs, sharp_edges=sharp_edges, record_history=record_history)
File "lightning-thunder/thunder/core/jit_ext.py", line 1660, in thunder_general_jit
process_recorded_modifications(ctx, epilogue_trace)
File "lightning-thunder/thunder/core/jit_ext.py", line 1554, in process_recorded_modifications
typ, name, root_module_provenance = get_parameter_or_buffer_or_submodule_name_and_root(
File "lightning-thunder/thunder/core/jit_ext.py", line 1255, in get_parameter_or_buffer_or_submodule_name_and_root
assert provenance.inputs[0].inst is PseudoInst.LOAD_ATTR
IndexError: list index out of range
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.