Lightning-AI / Lightning-AI/lightning-thunder
High Peak Memory with CUDAGraphTransform
@t-vi is already working on this.
Since Dec 11, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
Peak Memory is very high when CUDAGraphTransform is used.
```python
# With CUDAGraphTransform - 27517.101568
# Without CUDAGraphTransform - 11917.129728
```
Example -
```python
import torch
import thunder
import litgpt
from torch.testing import make_tensor
from functools import partial
from thunder.dynamo import ThunderCompiler
from thunder.transforms.cudagraph import CUDAGraphTransform
device = torch.device("cuda")
cfg = litgpt.Config.from_name("open_llama_3b", n_layer=10)
with device:
make = partial(make_tensor, low=0, high=255, device=device, dtype=torch.long, requires_grad=False)
shape = (1,) + (cfg.block_size,)
x = make(shape)
m = litgpt.GPT(cfg)
# m = thunder.jit(m)
m = thunder.jit(m, transforms=[CUDAGraphTransform()])
o = m(x)
o.sum().backward()
# With CUDAGraphTransform - 27517.101568
# Without CUDAGraphTransform - 11917.129728
print(torch.cuda.max_memory_allocated() / 1e6)
```
Tested with internal image dated `20241209` on RTX 6000 Ada.
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.