Lightning-AI / Lightning-AI/lightning-thunder
thunder may treat global (maybe nonlocal) value as constant in computation trace without a check in prologue
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
```python
import torch
import thunder
from contextvars import ContextVar
_compile_data = ContextVar("compile_data", default=1)
def fn(x):
v = _compile_data.get()
return x + v
jfn = thunder.jit(fn)
o = jfn(torch.ones(3,))
print(o) # tensor([2., 2., 2.])
_compile_data.set((2,))
o = jfn(torch.ones(3,))
print(o) # tensor([2., 2., 2.]) (should be tensor([3., 3., 3.]))
print(thunder.last_prologue_traces(jfn)[-1])
print(thunder.last_traces(jfn)[-1])
```
Prologue Trace
```python
@torch.no_grad()
@no_autocast
def prologue(*args, **kwargs):
# args: "Any"
check_len(args, 1)
# prims.check_len(args, 1)
# kwargs: "Any"
check_len(kwargs, 0)
# prims.check_len(kwargs, 0)
x: "cpu f32[3]" = args[0]
check_tensor_metadata(x, (3,), 'cpu', torch.float32, False)
# prims.check_tensor_shape_and_metadata(x, (3,), 'cpu', torch.float32, False)
cache_info: "Any" = thunder._get_cache_info()
cache_info_default_dtype: "" = cache_info['default_dtype']
check_literal_like(cache_info_default_dtype, torch.float32)
# prims.check_literal_like(cache_info_default_dtype, torch.float32)
cache_info_default_device: "" = cache_info['default_device']
check_literal_like(cache_info_default_device, torch.device("cpu"))
# prims.check_literal_like(cache_info_default_device, torch.device("cpu"))
cache_info_is_autocast_enabled: "bool False" = cache_info['is_autocast_enabled']
check_number_type_and_value(cache_info_is_autocast_enabled, False)
# prims.check_number_type_and_value(cache_info_is_autocast_enabled, False)
cache_info_no_grad_sync: "bool False" = cache_info['no_grad_sync']
check_number_type_and_value(cache_info_no_grad_sync, False)
# prims.check_number_type_and_value(cache_info_no_grad_sync, False)
cache_info_alias_tensor_indices: "str" = cache_info['alias_tensor_indices']
check_string_value(cache_info_alias_tensor_indices, '')
# prims.check_string_value(cache_info_alias_tensor_indices, '')
cache_info_is_grad_enabled: "bool True" = cache_info['is_grad_enabled']
check_number_type_and_value(cache_info_is_grad_enabled, True)
# prims.check_number_type_and_value(cache_info_is_grad_enabled, True)
return ((x,), ())
```
Computation Trace
```python
@torch.no_grad()
@no_autocast
def computation(x):
# x: "cpu f32[3]"
t0 = torch.add(x, 1, alpha=1) # t0: "cpu f32[3]"
# t0 = ltorch.add(x, 1, alpha=1) # t0: "cpu f32[3]"
# _ = prims.convert_element_type(1, float)
# t0 = prims.add(x, 1.0) # t0: "cpu f32[3]"
return t0
```
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 by running the provided ContextVar reproducer with thunder.jit and inspect the output from last_prologue_traces and last_traces. Trace how the value returned by _compile_data.get() is represented across compilation and where prologue checks are generated. Done means changing the ContextVar value produces tensor([3., 3., 3.]) rather than reusing the computation with 1.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100