Lightning-AI / Lightning-AI/lightning-thunder

have a method to compare speed of different parts of training between compilation backends

Open
#444 7 comments 2 reactions 1 assignee View on GitHub

@IvanYashchuk is already working on this.

Since May 28, 2024.

enhancement
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?

![Screenshot 2024-05-22 at 20 16 34](https://github.com/Lightning-AI/lightning-thunder/assets/149149379/979ca390-4a6e-4a19-827e-735b1a985b57)

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.