Lightning-AI / Lightning-AI/lightning-thunder
get_alloc_memory overestimates the peak allocated memory on Colab because Flash Attention backend is not chosen
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
`get_alloc_memory` has a potential to be a great tool for estimating the effect on memory usage after transforming traces. For those who don't know about this function it lives here:
https://github.com/Lightning-AI/lightning-thunder/blob/8c5905fd1a93145e690791a7c7a3c3e10b16b32b/thunder/examine/memory_caculation.py#L120-L137
I've tried generating the execution trace with Fake CUDA Tensors so that it's possible to analyze the final trace without the actual execution even on low-memory GPUs like in Colab:
```py
from litgpt.config import Config
from litgpt import GPT
import thunder
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
fake_mode = FakeTensorMode()
config = Config.from_name("Llama-3-8B")
with fake_mode, torch.device("cuda"):
model = GPT(config).to(torch.bfloat16)
with fake_mode:
micro_batch_size = 1
x = torch.randint(0, (micro_batch_size, model.max_seq_length,), device="cuda")
cmodel = thunder.jit(model)
cache_entry, inps, pro_to_epi = thunder.compile_data(cmodel).get_computation_and_inputs(x)
```
The estimated peak memory for this trace is for some reason different on Colab and locally
```py
import thunder.examine.memory_caculation
estimated_peak = thunder.examine.memory_caculation.get_alloc_memory(cache_entry.computation_traces[-1])[0]
print(f"Estimated peak memory usage in GiB: {estimated_peak / (1024 ** 3)}")
print(f"Estimated peak memory usage in GB: {estimated_peak / (1000 ** 3)}")
```
on Colab I see
```py
Estimated peak memory usage in GiB: 185.74466705322266
Estimated peak memory usage in GB: 199.4418176
```
and locally:
```py
Estimated peak memory usage in GiB: 56.76426076889038
Estimated peak memory usage in GB: 60.950160896
```
On Colab litgpt, nvfuser, and Thunder are installed with:
```py
!pip install litgpt
!pip install nvfuser-cu121-torch23
!pip install git+https://github.com/Lightning-AI/lightning-thunder@main
```
cc @apaz-cli
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with thunder/examine/memory_caculation.py, especially get_alloc_memory, and reproduce the reported trace using Fake CUDA Tensors and the Colab and local installation commands in the issue. Investigate why the Flash Attention backend is not chosen on Colab, then verify that the estimated peak memory is consistent with the selected backend.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100