Lightning-AI / Lightning-AI/lightning-thunder
improve shape accuracy in transform output by providing better tooling
Open
@t-vi is already working on this.
Since Sep 23, 2024.
enhancement
triage review
ux
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Some transforms, notably FSDP and TensorParallel ones, change shapes, but currently do not completely update them (it does for the linear that follows, but not for the activation etc.).
We might consider either making it easy for the transform to update the shapes on the fly or at the end of the transform by providing a mechanism similar / based on to the `interpret_trace_to_trace` introduced in #1164 (which would then become stricter by default).
cc: @crcrpar
cc @carmocca @borda
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.
Assessment
This issue has not been assessed yet.