Lightning-AI / Lightning-AI/lightning-thunder

Traces with bsyms of `torch.autograd.Function` are not printable

Open
#1,404 2 comments 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

*Note*: If you have a model or program that is not supported yet but should be, please use the program coverage template.

## 🐛 Bug

If traces include one or more `BoundSymbol`s of `torch.autograd.Function`, then they are not printable.

The lookaside is https://github.com/Lightning-AI/lightning-thunder/blob/3390c922cfbe1b70c42118dfa8aa71adb3bec692/thunder/core/jit_ext.py#L655.

When registering a symbol, it doesn't specify a module:https://github.com/Lightning-AI/lightning-thunder/blob/3390c922cfbe1b70c42118dfa8aa71adb3bec692/thunder/core/jit_ext.py#L688-L691, arriving at https://github.com/Lightning-AI/lightning-thunder/blob/3390c922cfbe1b70c42118dfa8aa71adb3bec692/thunder/core/symbol.py#L604-L616

### To Reproduce

#### Code sample

```python
import torch

import thunder

class MyLinear(torch.autograd.Function):
@staticmethod
def forward(ctx, x: torch.Tensor, weight: torch.Tensor, shape: tuple[int, int]) -> torch.Tensor:
ctx.shape = shape
ctx.save_for_backward(x, weight)
ctx.pretty_attr = 100
ctx.scaler = 1.0
return torch.matmul(x, weight.t())

@staticmethod
def backward(ctx, grad_output):
(x, weight) = ctx.saved_tensors
assert weight.shape == ctx.shape # really bogus, just to use ctx.shape
scaler2 = ctx.shape[0] / ctx.shape[1]
return torch.matmul(grad_output, weight) * ctx.scaler, torch.matmul(grad_output.t(), x) / scaler2, None

class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.l1 = torch.nn.Linear(2, 2, bias=False)

def forward(self, x):
return MyLinear.apply(x, self.l1.weight, self.l1.weight.shape)

if __name__ == "__main__":
x = torch.randn((2, 2), dtype=torch.float64, requires_grad=True)
model = Model().to(dtype=torch.float64)
jitted = thunder.jit(model)

jitted(x)
print(thunder.last_traces(jitted)[-1])
print(thunder.last_traces(jitted)[0])
```

The outputP:
```
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x, weight):
# x: "cpu f64[2, 2]"
# weight: "cpu f64[2, 2]"
t21 = torch.permute(weight, (1, 0)) # t21: "cpu f64[2, 2]"
# t21 = ltorch.permute(weight, (1, 0)) # t21: "cpu f64[2, 2]"
# t21 = prims.transpose(weight, (1, 0)) # t21: "cpu f64[2, 2]"
t22 = torch.matmul(x, t21) # t22: "cpu f64[2, 2]"
# t22 = ltorch.matmul(x, t21) # t22: "cpu f64[2, 2]"
# t22 = prims.matmul(x, t21) # t22: "cpu f64[2, 2]"
del t21
t11 = shallow_copy(t22) # t11: "cpu f64[2, 2]"
del t22
return {'output': t11, 'flat_args': [x, weight], 'flat_output': (t11,)}, ((weight, x), ())
Traceback (most recent call last):
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/repro.py", line 39, in
print(thunder.last_traces(jitted)[0])
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/core/trace.py", line 510, in __repr__
return self.python(print_depth=-1)
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/core/trace.py", line 363, in python
import_ctx, call_ctx, object_ctx = self._gather_ctxs()
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/core/trace.py", line 322, in _gather_ctxs
bsym_import_ctx, bsym_call_ctx, bsym_object_ctx = bsym.gather_ctxs()
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/core/symbol.py", line 648, in gather_ctxs
return self.import_ctx(), self._get_call_ctx(), self.object_ctx()
File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/core/symbol.py", line 608, in import_ctx
assert self.sym.module is not None # TODO: Is this a valid assumption?
AssertionError
```

### Expected behavior

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

Run the provided reproducer and inspect thunder/core/jit_ext.py around lines 655 and 688-691, then follow the symbol handling in thunder/core/symbol.py lines 604-616 and 648. Check the trace printing path in thunder/core/trace.py around lines 322 and 363; done means traces containing torch.autograd.Function BoundSymbols can be printed without the shown assertion.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.