Lightning-AI / Lightning-AI/lightning-thunder
[RFC] Option to make a trace easier to interpret
- 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