Lightning-AI / Lightning-AI/lightning-thunder
Multiple accesses for non-cached property fails in Thunder JIT
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
It seems that Thunder JIT currently assumes that attribute accesses always return the same object. It's true only for `@functools.cached_property`, if a class method is decorated with `@property` it should be treated as a method call.
```py
import torch
import thunder
class Test(torch.nn.Module):
@property
def test(self):
return object()
def forward(self):
return self.test, self.test
jtest = thunder.jit(Test())
jtest()
```
traceback:
```py
File ~/dev/lightning-thunder/thunder/core/interpreter.py:6164, in _call_dispatch(compilectx, runtimectx, fn, *args, **kwargs)
6162 if lookaside_fn:
6163 runtimectx.record_lookaside(lookaside_fn)
-> 6164 res = lookaside_fn(*args, **kwargs)
6165 return res
6167 # TODO: disabled as partial is just like any other class
6168 # (3) Handles partial objects
File ~/dev/lightning-thunder/thunder/core/jit_ext.py:732, in _general_jit_getattr_lookaside(obj, name, *maybe_default)
729 getattr_lookaside = default_lookaside(getattr)
730 assert getattr_lookaside is not None
--> 732 value = getattr_lookaside(obj, name, *maybe_default)
733 if value is INTERPRETER_SIGNALS.EXCEPTION_RAISED:
734 return value
File ~/dev/lightning-thunder/thunder/core/interpreter.py:1678, in _getattr_lookaside(obj, name, *maybe_default)
1676 if result is not INTERPRETER_SIGNALS.EXCEPTION_RAISED or not isinstance(ctx.curexc, AttributeError):
1677 if result is not INTERPRETER_SIGNALS.EXCEPTION_RAISED and compilectx._with_provenance_tracking:
-> 1678 result = wrap_attribute(result, obj, name)
1679 return result
1681 # `__getattr__` is only triggered if `__getattribute__` fails.
1682 # TODO: this should be `_interpret_call_with_unwrapping(getattr, obj, "__getattr__", null := object())`, but that would require multiple current exceptions.
File ~/dev/lightning-thunder/thunder/core/interpreter.py:1632, in wrap_attribute(plain_result, obj, name)
1629 # note: there are cases where "is" will always fail (e.g. BuiltinMethods
1630 # are recreated every time)
1631 if known_wrapper is not None:
-> 1632 assert plausibly_wrapper_of(
1633 known_wrapper, plain_result
1634 ), f"attribute {name.value} of {type(obj.value).__name__} object out of sync: {known_wrapper.value} vs. {plain_result}"
1635 return known_wrapper
1637 pr = ProvenanceRecord(PseudoInst.LOAD_ATTR, inputs=[obj.provenance, name.provenance])
AssertionError: attribute test of Test object out of sync: vs.
```
cc @apaz-cli
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in thunder/core/interpreter.py at wrap_attribute and _getattr_lookaside, then inspect thunder/core/jit_ext.py at _general_jit_getattr_lookaside. Reproduce the Test module and repeated @property accesses shown in the issue; done means Thunder JIT no longer asserts when the property returns a fresh object on each access.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100