Lightning-AI / Lightning-AI/lightning-thunder

Thunder saves too many tensors for backward for a Transformer's residual connection pattern

Open
#1,368 2 comments 0 reactions 2 assignees View on GitHub

Nobody has claimed this yet.

bug memory use rematerialization
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🐛 Bug

### To Reproduce

Here's a problematic pattern where Thunder's rematerialization algorithm is not effective:

```py
from torch.nn import Linear
import torch
import thunder

block_size = 16384
n_embd = 2560
intermediate_size = 10240
with torch.device("cuda"):
fc = Linear(n_embd, intermediate_size, bias=False).to(torch.float16)
proj = Linear(intermediate_size, n_embd, bias=False).to(torch.float16)
fc_weight = fc.weight
proj_weight = proj.weight

def mlp(x: torch.Tensor, fc_weight, proj_weight) -> torch.Tensor:
x = torch.nn.functional.linear(x, fc_weight)
x = torch.exp(x) #torch.nn.functional.gelu(x, approximate="none")
return torch.nn.functional.linear(x, proj_weight)

def f(x, fc_weight, proj_weight):
# start of transformer block
x_normed = torch.exp(x)
attention_output = x_normed

x = attention_output + x
nx = torch.exp(x)
x = mlp(nx, fc_weight, proj_weight) + x
# end of transformer block
x = torch.exp(x)
return x

tf = torch.compile(f)
jf = thunder.jit(f)

x = torch.randn(1, block_size, n_embd, device="cuda", requires_grad=True, dtype=torch.float16)

tout = tf(x, fc_weight, proj_weight)
print(f"{len(tout.grad_fn.saved_tensors)=}")
print(f"Saved tensors size torch.compile: {sum([t.numel() * t.element_size() for t in tout.grad_fn.saved_tensors if t is not None]) / 2**20:.2f} MiB")

jout = jf(x, fc_weight, proj_weight)
print(f"{len(jout.grad_fn.saved_tensors)=}")
print(f"Saved tensors size Thunder: {sum([t.numel() * t.element_size() for t in jout.grad_fn.saved_tensors if t is not None]) / 2**20:.2f} MiB")
```
Script output:
```
len(tout.grad_fn.saved_tensors)=6
Saved tensors size torch.compile: 660.00 MiB
len(jout.grad_fn.saved_tensors)=8
Saved tensors size Thunder: 1060.00 MiB
```

Current `thunder.examine.make_trace_dot` doesn't print out all output variables so it's difficult to read the graph, but, anyway, here is the joint forward and backward graph before rematerialization is applied:
![image](https://github.com/user-attachments/assets/4de13555-544e-434d-8a84-df0a794c2a43)

Created with a breakpoint before this line: https://github.com/Lightning-AI/lightning-thunder/blob/9c916d9df73f3920b51e5951303a76b25ab2d4d4/thunder/core/rematerialization.py#L609

Resolving this bug would help us resolve https://github.com/Lightning-AI/lightning-thunder/issues/246.

cc @riccardofelluga

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.