Lightning-AI / Lightning-AI/lightning-thunder
Support FSDP and torch.compile
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
### Motivation
We use `async_op=True` to reorder the collective calls in the trace, but that's not supported by Dynamo:
```python
[rank7]: torch._dynamo.exc.Unsupported: CollectiveFunctionRewriteVariable can't support async_op=True for
[rank7]: from user code:
[rank7]: File "thunder.torch_interpreted_func_94", line 13, in torch_interpreted_func
[rank7]: p8 = torch_all_gather_prim_impl(t2, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p8: "FUTURE cuda:7 bf16[2560, 2048]"
[rank7]: File "/home/carlos/lightning-thunder/thunder/executors/torchex.py", line 1687, in _all_gather_prim_impl
[rank7]: handle: None | torch.distributed.distributed_c10d.Work = torch.distributed.all_gather_into_tensor(
```
### Pitch
We could support this two ways:
a. Disable `fullgraph=True` to allow graph-breaks via #281
b. Do not let the torch_compile_ex to fuse through these collective calls.
Option b) is what I would recommend.
### Alternatives
Do not support this.
### Additional context
Requires #140 to land first.
cc @carmocca @awaelchli @crcrpar @apaz-cli
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 with thunder/executors/torchex.py around the _all_gather_prim_impl entry point and review prerequisite issue #140. Compare the proposed graph-break and non-fusion approaches for async collective calls, then identify the FSDP and torch.compile scenarios that would demonstrate completion.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 28/100