Lightning-AI / Lightning-AI/lightning-thunder

Task list to jit `Adam.step`

Open
#1,055 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

optimizer program-coverage
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

  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.

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.