Metadata-Version: 2.4
Name: torch2-look-ahead
Version: 0.1.0
Summary: A modern PyTorch 2.x implementation of the Lookahead optimizer wrapper.
Author: Michael Zhang, James Lucas, Geoffrey Hinton, Jimmy Ba
Maintainer: Wintoplay
License-Expression: MIT
Project-URL: Homepage, https://github.com/Wintoplay/torch2-look-ahead
Project-URL: Repository, https://github.com/Wintoplay/torch2-look-ahead
Project-URL: Issues, https://github.com/Wintoplay/torch2-look-ahead/issues
Keywords: pytorch,optimizer,lookahead,deep-learning
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Provides-Extra: cifar
Requires-Dist: numpy; extra == "cifar"
Requires-Dist: torchvision; extra == "cifar"
Provides-Extra: dev
Requires-Dist: build; extra == "dev"
Requires-Dist: pytest; extra == "dev"
Requires-Dist: twine; extra == "dev"
Dynamic: license-file

# Lookahead Optimizer for PyTorch 2.x

This repository provides a modern PyTorch implementation of the Lookahead optimizer wrapper from [Lookahead Optimizer: k steps forward, 1 step back](https://arxiv.org/abs/1907.08610).

The PyTorch wrapper has been renewed for PyTorch 2.x optimizer behavior: schedulers can wrap it, `zero_grad(set_to_none=...)` is supported, checkpoint state includes both the inner optimizer and Lookahead slow weights, and slow-weight evaluation has public helper methods.

The TensorFlow file is kept as historical reference. The maintained path in this repo is `lookahead_pytorch.py`.

## Installation

Install from PyPI:

```bash
pip install torch2-look-ahead
```

Core optimizer usage only needs PyTorch, which is installed as a package dependency.

The CIFAR-10 comparison script also needs `torchvision` and `numpy`:

```bash
pip install torch torchvision numpy
```

Muon comparisons use `torch.optim.Muon`. If your PyTorch build does not expose `torch.optim.Muon`, the script will fail early with a clear message. PyTorch 2.0-era installs may not include it.

## Basic Usage

Wrap any PyTorch optimizer:

```python
import torch
from lookahead_pytorch import Lookahead

model = ...
base_optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
optimizer = Lookahead(base_optimizer, la_steps=5, la_alpha=0.8)

for inputs, targets in train_loader:
    optimizer.zero_grad(set_to_none=True)
    loss = loss_fn(model(inputs), targets)
    loss.backward()
    optimizer.step()
```

Schedulers should receive the Lookahead wrapper, not the inner optimizer:

```python
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=1e-3,
    epochs=epochs,
    steps_per_epoch=len(train_loader),
)
```

## Slow-Weight Evaluation

Lookahead often validates better with the slow weights. Use the public backup helpers to evaluate slow weights without permanently replacing the fast training weights:

```python
optimizer.backup_and_load_cache()
try:
    val_loss, val_acc = evaluate(model, valid_loader)
finally:
    optimizer.clear_and_load_backup()
```

The old private names `_backup_and_load_cache()` and `_clear_and_load_backup()` are still available as compatibility aliases.

## Checkpointing

Save the model state and the Lookahead optimizer state together:

```python
torch.save(
    {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
    },
    "checkpoint.pt",
)
```

Restore by constructing the same model and inner optimizer first, then loading both states:

```python
checkpoint = torch.load("checkpoint.pt", map_location="cpu")
model.load_state_dict(checkpoint["model"])

base_optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
optimizer = Lookahead(base_optimizer, la_steps=5, la_alpha=0.8)
optimizer.load_state_dict(checkpoint["optimizer"])
```

Old checkpoints created by the original wrapper, where `Lookahead.state_dict()` returned only the inner optimizer state, still load. In that case the slow weights are initialized from the current model parameters.

## Momentum Pullback

`pullback_momentum` supports the original options: `none`, `reset`, and `pullback`.

Momentum pullback only applies when the inner optimizer stores a `momentum_buffer` for a parameter. Optimizers such as AdamW do not use that buffer, so the option is skipped safely for those parameters.

## CIFAR-10 ResNet9 Comparison

`compare_cifar10_optimizers.py` is a standalone script based on `05b_cifar10_resnet.ipynb`. It trains the notebook's ResNet9 model with CIFAR-10 normalization, random crop/flip augmentation, gradient clipping, weight decay, and one-cycle learning-rate scheduling. Momentum cycling is disabled so the same scheduler works for both single-optimizer and mixed Muon+AdamW runs. AdamW uses conservative `foreach=False` and `fused=False` defaults because those are more reliable on ROCm builds.

Run a fast smoke comparison:

```bash
python compare_cifar10_optimizers.py --preset smoke
```

Run the full notebook-like comparison:

```bash
python compare_cifar10_optimizers.py --preset full
```

Compare a subset of optimizers:

```bash
python compare_cifar10_optimizers.py --preset smoke --optimizers adamw,lookahead_adamw
python compare_cifar10_optimizers.py --preset smoke --optimizers muon,lookahead_muon
```

Available optimizer names:

- `adamw`
- `muon`
- `lookahead_adamw`
- `lookahead_muon`

The script writes summary JSON metrics to `results/cifar10_optimizer_comparison.json` by default. Normal runs isolate each optimizer in its own child process and stream completed epochs to `results/cifar10_optimizer_comparison.jsonl`, so a ROCm abort in one optimizer does not erase earlier metrics or stop later optimizers. Use `--no-isolate-optimizers` only when debugging a raw single-process failure.

Important Muon detail: this PyTorch implementation of `torch.optim.Muon` only accepts 2-D parameters. ResNet9 contains convolution, batch-normalization, and bias tensors that are not 2-D, so the script uses Muon for eligible 2-D matrix parameters and AdamW as a fallback for the rest. `lookahead_muon` wraps that combined update, so Lookahead slow weights cover the whole model.

Useful options:

```bash
python compare_cifar10_optimizers.py \
  --preset smoke \
  --optimizers adamw,muon,lookahead_adamw,lookahead_muon \
  --max-lr 0.01 \
  --weight-decay 1e-4 \
  --grad-clip 0.1 \
  --lookahead-steps 5 \
  --lookahead-alpha 0.8 \
  --seed 42
```

Preset defaults:

| Preset | Epochs | Batch size | Train examples | Validation examples | Workers |
| --- | ---: | ---: | ---: | ---: | ---: |
| `smoke` | 1 | 128 | 2048 | 1024 | 0 |
| `full` | 8 | 400 | full CIFAR-10 train split | full CIFAR-10 test split | 0 |

Use `--train-subset 0 --val-subset 0` to force the full dataset with either preset, or pass explicit subset sizes.

`num_workers` controls PyTorch DataLoader worker processes. The full preset uses `0` because this ROCm setup reproduced HIP illegal-memory-access failures with multi-worker loading and completed normally with single-process loading.

### ROCm / CUDA launch failures

By default, `--device auto` runs a small CUDA probe in a subprocess before training. If the selected optimizer path crashes the accelerator backend, the script prints the probe output and falls back to CPU instead of aborting the main run.

If you see a HIP launch failure like `hipErrorLaunchFailure`, use one of these commands:

```bash
python compare_cifar10_optimizers.py --preset smoke --device cpu
python compare_cifar10_optimizers.py --preset smoke --batch-size 32
AMD_SERIALIZE_KERNEL=3 python compare_cifar10_optimizers.py --preset smoke --device cuda --no-cuda-probe
```

AdamW backend flags are exposed for debugging:

```bash
python compare_cifar10_optimizers.py --preset smoke --adamw-foreach false --adamw-fused false
python compare_cifar10_optimizers.py --preset smoke --adamw-foreach auto --adamw-fused auto
```

The first command is the safer ROCm default. The second lets PyTorch choose its faster backend and may reproduce backend-specific crashes.

## Development Checks

Run the focused tests:

```bash
python -m pytest tests/test_lookahead_pytorch.py
```

Compile the script and optimizer:

```bash
python -m py_compile lookahead_pytorch.py compare_cifar10_optimizers.py tests/test_lookahead_pytorch.py
```

## Citation

```bibtex
@article{zhang2019lookahead,
  title={Lookahead Optimizer: k steps forward, 1 step back},
  author={Zhang, Michael R and Lucas, James and Hinton, Geoffrey and Ba, Jimmy},
  journal={arXiv preprint arXiv:1907.08610},
  year={2019}
}
```

<img src="figs/accuracy_surface.png" width="500">
