Lightning-AI / Lightning-AI/lightning-thunder

get_alloc_memory overestimates the peak allocated memory on Colab because Flash Attention backend is not chosen

Open
#677 7 comments 0 reactions 0 assignees View on GitHub
examine memory use
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.