Lightning-AI / Lightning-AI/lightning-thunder
have a method to compare speed of different parts of training between compilation backends
@IvanYashchuk is already working on this.
Since May 28, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Have a method to annotate pieces of training code (e.g. [benchmark_litgpt](https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/benchmarks/benchmark_litgpt.py)) so we can easily and automatically compare effectiveness of different compilation methods / versions of Thunder on these pieces instead of whole training loop.
### Motivation
Right now we provide regular benchmarking of Thunder on LitGPT models. If we spot that Thunder is slower than Inductor or there is a regression, the cause must be found manually by investigating logs / mapping kernels from nsys profiling reports, which have different names. It's not easy to speed-up and automate this process.
If we were able to split the training loop into smaller chunks and annotate them, we could automatically find the chunks that are slower / can be improved and provide them together with benchmarking results. This would allow us to spot issues to be improved much faster.
### Pitch
One idea would be to use NXTV markers - I was able to add them to LitGPT models, so we know the duration of forward passes for each module.
It seems to look good for Inductor and Eager mode, however for Thunder some of the markers got removed. Maybe there would be a way to preserve them with some rules applied if e.g. there is problem with fusing operations?

The problem also is that in this way we miss the backward passes. I found that it can be solved by using: https://pytorch.org/docs/stable/autograd.html#torch.autograd.profiler.emit_nvtx
and I'm looking into it right now.
### Alternatives
If there are some alternatives, please let me know.
### Additional context
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.