Lightning-AI / Lightning-AI/lightning-thunder

[RFC] Option to make a trace easier to interpret

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

Description

## 🚀 Feature

When we run the following snippet with `torchrun --nproc-per-node=2 sample.py`, we'd get the trace below, as of c89fb1bc5114161aed78e9eb19b8f68df523d192.

```python
import os

import torch
from torch.distributed import distributed_c10d as c10d

import thunder
from thunder.tests.distributed.helper import ToyModel
from thunder.distributed import column_parallel, row_parallel

def main():
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device(f"cuda:{local_rank}")
c10d.init_process_group()

model = ToyModel(bias=True).to(device)
if local_rank == 0:
print("# Convert `net1` and `net2` into column-wise parallel and row-wise parallel, respectively")
print(model)
tp_model = thunder.jit(model)
tp_model = column_parallel(tp_model, ["net1"])
tp_model = row_parallel(tp_model, ["net2"])

x = torch.randn((1, ToyModel.N_IN), device=device)
tp_model(x)

if local_rank == 0:
fw_extrace = thunder.last_traces(tp_model)[-1]

print(fw_extrace)

c10d.destroy_process_group()

if __name__ == "__main__":
main()
```

```python
# Convert `net1` and `net2` into column-wise parallel and row-wise parallel, respectively
# ToyModel(
# (net1): Linear(in_features=12, out_features=16, bias=True)
# (net2): Linear(in_features=16, out_features=8, bias=True)
# )

# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from torch import Tensor
import torch.nn.functional
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def augmented_forward_fn(x, t_net1_bias, t_net1_weight, t_net2_bias, t_net2_weight):
# x: "cuda:0 f32[1, 12]"
# t_net1_bias: "cuda:0 f32[8]"
# t_net1_weight: "cuda:0 f32[8, 12]"
# t_net2_bias: "cuda:0 f32[8]"
# t_net2_weight: "cuda:0 f32[8, 8]"
t0 = torch.nn.functional.linear(x, t_net1_weight, t_net1_bias) # t0: "cuda:0 f32[1, 8]"
# t0 = ltorch.linear(x, t_net1_weight, t_net1_bias) # t0: "cuda:0 f32[1, 8]"
# t0 = prims.linear(x, t_net1_weight, t_net1_bias) # t0: "cuda:0 f32[1, 8]"
p1 = torch_all_gather_prim_impl(t0, _torch_distributed_distributed_c10d_ProcessGroup_0, True, 0) # p1: "FUTURE cuda:0 f32[2, 8]"
del t0
t2 = torch_wait_prim_impl(p1) # t2: "cuda:0 f32[2, 8]"
del p1
(t3, t4) = torch.chunk(t2, 2, 0)
# (t3, t4) = ltorch.chunk(t2, 2, 0)
# t3 = prims.slice_prim(t2, [0, 0], [1, 8], [1, 1]) # t3: "cuda:0 f32[1, 8]"
# t4 = prims.slice_prim(t2, [1, 0], [2, 8], [1, 1]) # t4: "cuda:0 f32[1, 8]"
del t2
[t13, t5] = nvFusion0(t3, t4)
# t5 = prims.cat((t3, t4), -1) # t5: "cuda:0 f32[1, 16]"
# t6 = prims.mul(0.5, t5) # t6: "cuda:0 f32[1, 16]"
# t7 = prims.pow(t5, 3.0) # t7: "cuda:0 f32[1, 16]"
# t8 = prims.mul(0.044715, t7) # t8: "cuda:0 f32[1, 16]"
# t9 = prims.add(t5, t8) # t9: "cuda:0 f32[1, 16]"
# t10 = prims.mul(0.7978845608028654, t9) # t10: "cuda:0 f32[1, 16]"
# t11 = prims.tanh(t10) # t11: "cuda:0 f32[1, 16]"
# t12 = prims.add(1.0, t11) # t12: "cuda:0 f32[1, 16]"
# t13 = prims.mul(t6, t12) # t13: "cuda:0 f32[1, 16]"
del t3, t4
t14 = torch_slice_prim_impl(t13, [0, 0], [1, 8], [1, 1]) # t14: "cuda:0 f32[1, 8]"
del t13
t15 = torch.nn.functional.linear(t14, t_net2_weight, None) # t15: "cuda:0 f32[1, 8]"
# t15 = ltorch.linear(t14, t_net2_weight, None) # t15: "cuda:0 f32[1, 8]"
# t15 = prims.linear(t14, t_net2_weight, None) # t15: "cuda:0 f32[1, 8]"
p16 = torch_all_reduce_prim_impl(t15, _DistributedReduceOps_1, _torch_distributed_distributed_c10d_ProcessGroup_0, True, True) # p16: "FUTURE cuda:0 f32[1, 8]"
del t15
t17 = torch_wait_prim_impl(p16) # t17: "cuda:0 f32[1, 8]"
del p16
t29 = torch.unsqueeze(t_net2_bias, 0) # t29: "cuda:0 f32[1, 8]"
# t29 = ltorch.unsqueeze(t_net2_bias, 0) # t29: "cuda:0 f32[1, 8]"
# t29 = prims.broadcast_in_dim(t_net2_bias, [1, 8], [1]) # t29: "cuda:0 f32[1, 8]"
t18 = Tensor.expand(t29, (1, 8)) # t18: "cuda:0 f32[1, 8]"
# t18 = ltorch.expand(t29, (1, 8)) # t18: "cuda:0 f32[1, 8]"
# t18 = prims.broadcast_in_dim(t29, (1, 8), (0, 1)) # t18: "cuda:0 f32[1, 8]"
del t29
[t19] = nvFusion1(t17, t18)
# t19 = prims.add(t17, t18) # t19: "cuda:0 f32[1, 8]"
del t17, t18
return {'output': t19, 'flat_args': [x, t_net1_bias, t_net1_weight, t_net2_bias, t_net2_weight], 'flat_output': (t19,)}, ((t14, t5, t_net2_weight, x), (0.5, 3.0, 0.044715, 0.7978845608028654))
```

In this trace, some `BoundSymbol`s are derived from tensor parallel pre/post-processing. In my humble opinion the trace would look easier to interpret if the corresponding bsyms have a label or a name indicating they are pre/post-processing.

One example would be:
```python
p16 = torch_all_reduce_prim_impl(t15, _DistributedReduceOps_1, _torch_distributed_distributed_c10d_ProcessGroup_0, True, True) # p16: "FUTURE cuda:0 f32[1, 8]"
del t15
t17 = torch_wait_prim_impl(p16) # t17: "cuda:0 f32[1, 8]"
del p16
t29 = torch.unsqueeze(t_net2_bias, 0) # t29: "cuda:0 f32[1, 8]"
# t29 = ltorch.unsqueeze(t_net2_bias, 0) # t29: "cuda:0 f32[1, 8]"
# t29 = prims.broadcast_in_dim(t_net2_bias, [1, 8], [1]) # t29: "cuda:0 f32[1, 8]"
t18 = Tensor.expand(t29, (1, 8)) # t18: "cuda:0 f32[1, 8]"
# t18 = ltorch.expand(t29, (1, 8)) # t18: "cuda:0 f32[1, 8]"
# t18 = prims.broadcast_in_dim(t29, (1, 8), (0, 1)) # t18: "cuda:0 f32[1, 8]"
del t29
[t19] = nvFusion1(t17, t18)
# t19 = prims.add(t17, t18) # t19: "cuda:0 f32[1, 8]"
```

If these lines are easy to guess they are post-processing of row-wise parallel linear, I'd say the trace is more perspicuous.
An alternative is that we call a method named `row_linear_postprocessing` in a visitor_transform and a bsym would just call it.

### Motivation

I've had experienced an appetite that I'd like to have a `BoundSymbol` call a sequence of primitive ops without decomposition / executor-specific optimizations, mainly for the sake of readability/clarity of it in a trace.

### Pitch

### Alternatives

### Additional context

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the thunder.last_traces(tp_model) entry point and the tensor-parallel visitor_transform mentioned in the issue, then trace how BoundSymbols are represented and rendered. Clarify and implement a design for optionally labeling tensor-parallel pre/post-processing sequences, with the trace becoming easier to interpret without decomposition or executor-specific optimizations.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, distributed-systems
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.