Lightning-AI / Lightning-AI/lightning-thunder

Updating a nn.Module attribute in forward raises an exception in prologue trace.

Open
#137 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug jit warnings & errors
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

```python
import torch
import thunder

import thunder.examine

class MyModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.bar = 1

def forward(self, x):
self.bar = self.bar + 1
# self.bar = 2 # This works
return x

m = MyModule()

x = torch.randn(16, 16, device='cuda')

jit_linear = thunder.jit(m)

o = jit_linear(x)
```

Error:
```
File "/home/kkalambarkar/lightning-thunder/thunder/__init__.py", line 537, in get_computation_and_inputs
inps = pro(*args, **kwargs)
File "/home/kkalambarkar/git/pytorch/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/home/kkalambarkar/miniconda3/envs/pytorch-dev/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "thunder.prologue_0", line 16, in prologue
File "/home/kkalambarkar/lightning-thunder/thunder/executors/pythonex.py", line 100, in _check_number_type_and_value_impl
utils.check(
File "/home/kkalambarkar/lightning-thunder/thunder/core/baseutils.py", line 103, in check
raise exception_type(s())
RuntimeError: Expected 2 to be equal to and have the type of 1
```

cc @apaz-cli

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 by reproducing the example in the issue, then inspect thunder/__init__.py around get_computation_and_inputs and thunder/executors/pythonex.py around _check_number_type_and_value_impl. Trace how the prologue handles self.bar changing from 1 to 2. Done means the example no longer raises the type/value error while preserving the attribute update.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.