Lightning-AI / Lightning-AI/lightning-thunder
Implement gelu (and other elementwise fusions) recomputation during backward
@riccardofelluga is already working on this.
Since Aug 21, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
### Motivation
```py
import torch
import thunder
def f(x):
x = x @ x
x = torch.nn.functional.gelu(x, approximate="none")
x = x @ x
return x
x = torch.randn(1000, 1000, device="cuda", requires_grad=True)
jf = thunder.jit(f)
y = jf(x)
print([t.shape for t in y.grad_fn.saved_tensors])
# This will print:
# [torch.Size([1000, 1000]), torch.Size([1000, 1000]), torch.Size([1000, 1000])]
print(thunder.core.vjp_utils.get_saved_for_backward_tensors(thunder.last_traces(jf)[-1]))
# This will print:
# (,
# ,
# )
```
In the above snippet `t5` is the output of the gelu function and the request is to implement a pass that forces recomputation of the gelu function in the backward pass instead of saving this intermediate tensor.
Ongoing PR: https://github.com/Lightning-AI/lightning-thunder/pull/1003.
Implementing gelu recomputation would resolve the OOM error seen in https://github.com/Lightning-AI/lightning-thunder/issues/246.
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.