Lightning-AI / Lightning-AI/lightning-thunder

Build op provenance tracking into compile trace output

Open
#255 0 comments 3 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature

The request is to be able to connect the practitioner's model code clearly to the produced graph trace by Thunder. Ideally, each traced node should be able to map back to the model code for which it got generated.

### Motivation

This would tremendously help debugging issues around graph capture, graph optimization (such as rematerialization, DCE etc.). This also helps improve user understanding of what Thunder is doing. It could also be very helpful for developers to build tools that can operate on top of Thunder graphs.
Examples from TorchInductor in the Pitch.

### Pitch/Additional Context

Example of FX Graph debug from TorchInductor - mapping traced graph decomposed nodes back to practitioner model code.
```
# File: /scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py:35, code: out = self.conv(x)
convolution: f32[16, 64, 56, 56] = torch.ops.aten.convolution.default(primals_7, primals_1, None, [1, 1], [0, 0], [1, 1], False, [0, 0], 1)

# File: /scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py:36, code: out = self.bn(out)
add: i64[] = torch.ops.aten.add.Tensor(primals_6, 1); primals_6 = None
var_mean = torch.ops.aten.var_mean.correction(convolution, [0, 2, 3], correction = 0, keepdim = True)
getitem: f32[1, 64, 1, 1] = var_mean[0]
getitem_1: f32[1, 64, 1, 1] = var_mean[1]; var_mean = None
add_1: f32[1, 64, 1, 1] = torch.ops.aten.add.Tensor(getitem, 1e-05)
rsqrt: f32[1, 64, 1, 1] = torch.ops.aten.rsqrt.default(add_1); add_1 = None
sub: f32[16, 64, 56, 56] = torch.ops.aten.sub.Tensor(convolution, getitem_1)
mul: f32[16, 64, 56, 56] = torch.ops.aten.mul.Tensor(sub, rsqrt); sub = None
squeeze: f32[64] = torch.ops.aten.squeeze.dims(getitem_1, [0, 2, 3]); getitem_1 = None
squeeze_1: f32[64] = torch.ops.aten.squeeze.dims(rsqrt, [0, 2, 3]); rsqrt = None
mul_1: f32[64] = torch.ops.aten.mul.Tensor(squeeze, 0.1)
mul_2: f32[64] = torch.ops.aten.mul.Tensor(primals_4, 0.9); primals_4 = None
add_2: f32[64] = torch.ops.aten.add.Tensor(mul_1, mul_2); mul_1 = mul_2 = None
squeeze_2: f32[64] = torch.ops.aten.squeeze.dims(getitem, [0, 2, 3]); getitem = None
mul_3: f32[64] = torch.ops.aten.mul.Tensor(squeeze_2, 1.0000199302441455); squeeze_2 = None
mul_4: f32[64] = torch.ops.aten.mul.Tensor(mul_3, 0.1); mul_3 = None
mul_5: f32[64] = torch.ops.aten.mul.Tensor(primals_5, 0.9); primals_5 = None
add_3: f32[64] = torch.ops.aten.add.Tensor(mul_4, mul_5); mul_4 = mul_5 = None
unsqueeze: f32[64, 1] = torch.ops.aten.unsqueeze.default(primals_2, -1)
unsqueeze_1: f32[64, 1, 1] = torch.ops.aten.unsqueeze.default(unsqueeze, -1); unsqueeze = None
unsqueeze_2: f32[64, 1] = torch.ops.aten.unsqueeze.default(primals_3, -1); primals_3 = None
unsqueeze_3: f32[64, 1, 1] = torch.ops.aten.unsqueeze.default(unsqueeze_2, -1); unsqueeze_2 = None
mul_6: f32[16, 64, 56, 56] = torch.ops.aten.mul.Tensor(mul, unsqueeze_1); mul = unsqueeze_1 = None
add_4: f32[16, 64, 56, 56] = torch.ops.aten.add.Tensor(mul_6, unsqueeze_3); mul_6 = unsqueeze_3 = None

# File: /scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py:37, code: out = self.relu(out)
relu: f32[16, 64, 56, 56] = torch.ops.aten.relu.default(add_4); add_4 = None
le: b8[16, 64, 56, 56] = torch.ops.aten.le.Scalar(relu, 0)
```
Similarly, in the final codegen output, one can see which decomposed node belongs in each kernel generated and what aten level op did the decomposed node come from. Inside the kernel, there are also comments describing the practitioner code stack which is included in each kernel.
*This is already covered by Thunder today to some extent* as the trace output lists all the decomposed nodes which are part of a NVFusion block. But a mapping to original code would be fantastic for better understanding.

```
# aten._native_batch_norm_legit_functional => add_1, add_4, mul, mul_6, rsqrt, sub, var_mean
# aten.relu => relu
# aten.threshold_backward => le
triton_poi_fused__native_batch_norm_legit_functional_relu_threshold_backward_4 = async_compile.triton('''
import triton
import triton.language as tl
from torch._inductor.ir import ReductionHint
from torch._inductor.ir import TileHint
from torch._inductor.triton_heuristics import pointwise
from torch._inductor.utils import instance_descriptor

@pointwise(size_hints=[4194304], filename=__file__, meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: '*fp32', 5: '*fp32', 6: '*i1', 7: 'i32'}, 'device': 0, 'constants': {}, 'mutated_arg_names': [], 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2, 3, 4, 5, 6, 7), equal_to_1=())]})
@triton.jit
def triton_poi_fused__native_batch_norm_legit_functional_relu_threshold_backward_4(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, out_ptr0, out_ptr1, xnumel, XBLOCK : tl.constexpr):
xnumel = 3211264
xoffset = tl.program_id(0) * XBLOCK
xindex = xoffset + tl.arange(0, XBLOCK)[:]
xmask = xindex < xnumel
x3 = xindex
x1 = (xindex // 3136) % 64

# ORIGIN:
# call_function aten.relu.default
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 37, in forward\ out = self.relu(out)\
# END ORIGIN

# ORIGIN:
# call_function aten.add.Tensor
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 36, in forward\ out = self.bn(out)\
# END ORIGIN

# ORIGIN:
# call_function aten.rsqrt.default
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 36, in forward\ out = self.bn(out)\
# END ORIGIN

# ORIGIN:
# call_function aten.mul.Tensor
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 36, in forward\ out = self.bn(out)\
# END ORIGIN

# ORIGIN:
# call_function aten.var_mean.correction
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 36, in forward\ out = self.bn(out)\
# END ORIGIN

# ORIGIN:
# call_function aten.mul.Tensor
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 36, in forward\ out = self.bn(out)\
# END ORIGIN

# ORIGIN:
# call_function aten.add.Tensor
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 36, in forward\ out = self.bn(out)\
# END ORIGIN

# ORIGIN:
# call_function aten.sub.Tensor
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 36, in forward\ out = self.bn(out)\
# END ORIGIN

tmp0 = tl.load(in_ptr0 + (x3), None)
tmp1 = tl.load(in_ptr1 + (x1), None)
tmp3 = tl.load(in_ptr2 + (x1), None)
tmp10 = tl.load(in_ptr3 + (x1), None)
tmp12 = tl.load(in_ptr4 + (x1), None)
tmp2 = tmp0 - tmp1
tmp4 = 50176.0
tmp5 = tmp3 / tmp4
tmp6 = 1e-05
tmp7 = tmp5 + tmp6
tmp8 = tl.math.rsqrt(tmp7)
tmp9 = tmp2 * tmp8
tmp11 = tmp9 * tmp10
tmp13 = tmp11 + tmp12
tmp14 = tl.where(0 != 0, 0, tl.where(0 > tmp13, 0, tmp13))

# ORIGIN:
# call_function aten.le.Scalar
# File "/scratch/mojitos/Pytorch/resnet/test_conv_bn_relu.py", line 37, in forward\ out = self.relu(out)\
# END ORIGIN

tmp15 = 0.0
tmp16 = tmp14 <= tmp15
tl.store(out_ptr0 + (x3 + tl.zeros([XBLOCK], tl.int32)), tmp14, None)
tl.store(out_ptr1 + (x3 + tl.zeros([XBLOCK], tl.int32)), tmp16, None)
''')
```

cc @carmocca

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by locating Thunder's compile trace output and final codegen paths, then read how traced and decomposed nodes are represented. The issue provides TorchInductor examples but names no Thunder files or tests. Done means trace and generated-kernel output can map nodes back to practitioner model code, with coverage for graph transformations and fused kernels.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.