torch2-look-ahead 0.1.0

A modern PyTorch 2.x implementation of the Lookahead optimizer wrapper.

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.

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:

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:

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:

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:

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:

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:

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:

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:

python compare_cifar10_optimizers.py --preset smoke

Run the full notebook-like comparison:

python compare_cifar10_optimizers.py --preset full

Compare a subset of optimizers:

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:

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:

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:

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:

python -m pytest tests/test_lookahead_pytorch.py

Compile the script and optimizer:

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

Citation

@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}
}