Lightning-AI / Lightning-AI/lightning-thunder
Task list to jit `Adam.step`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
This is a list of tasks to enable compiling `Adam.step`/`AdamW.step`.
## Goal
The expected usage is:
```py
optimizer = Adam(model.parameters())
optimizer_step = jit(optimizer.step)
# or torch.compile(optimizer.step, backend=)
for X, y in data:
optimizer.zero_grad()
outputs = model(X)
loss = loss_fn(X, y)
loss.backward()
optimizer_step()
```
## Task list for `torch.compile(optimizer.step, backend=)`
We have to be able to trace through PyTorch's implementation of `Adam.step` (with `foreach=False`, `fused=False`)
- [x] #1070
Turned out to be not necessary
- [x] #1062
## Task list for `thunder.jit(optimizer.step)`
~This can be easily achieved by applying an `OperatorExecutor` to symbolically replace `Adam.step`'s inner implementation with our own thunder-friendly implementation.~
This is no longer needed thanks to improved support for in-place ops.
We need the following to trace `Adam.step`:
* `_single_tensor_adam` function
- [ ] Add `Tensor.conj`
- [ ] Cast the computation result of the in-place operation into the dst tensor's dtype
Specifically, we must support `(bloat16[N]).addcdiv_(bfloat16[N], bfloat16[N], value=float32)`
This was done in [#1194](https://github.com/Lightning-AI/lightning-thunder/pull/1194)
* Inner state tensor management
- [ ] Add `memory_format` option to `zeros_like`
- [ ] #1057
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.
Research direction
Start by reading PyTorch's `_single_tensor_adam` implementation and the issue linked as #1057, then inspect the existing support for `Tensor.conj`, dtype casting in `addcdiv_`, and `zeros_like(memory_format)`. Done means the listed tracing requirements are implemented so `thunder.jit(optimizer.step)` can compile both Adam and AdamW step paths.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100