Lightning-AI / Lightning-AI/lightning-thunder
Thunder saves too many tensors for backward for a Transformer's residual connection pattern
Nobody has claimed this yet.
- 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:

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
- 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.