Schedule-Free Optimization in PyTorch

Common errors

• Exception: Optimizer was not in train mode when step is called. optimizer.train() must be called before optimizer.step(). See documentation for details.

  cause The `optimizer.train()` method was not called before `optimizer.step()` during the training loop.
  fix Add `optimizer.train()` at the beginning of your training epoch or step, similar to how `model.train()` is used.

• Incorrect or inconsistent evaluation results, especially with models using BatchNorm layers.

  cause Batch normalization statistics during evaluation are being computed from the intermediate `y` sequence of the optimizer instead of the final `x` sequence, leading to discrepancies.
  fix Implement explicit handling for BatchNorm layers during evaluation to ensure statistics are correctly updated from the `x` sequence, or use `PreciseBN` if applicable.

• Suboptimal convergence or performance compared to traditional optimizers with well-tuned learning rate schedules.

  cause Despite being 'schedule-free', the library still requires careful tuning of other hyperparameters like the initial learning rate and regularization, and potentially the `beta` parameter. Starting with default values without adjustment may yield poor results.
  fix Tune the learning rate (often higher than traditional optimizers, e.g., 10x-50x for SGD, 1x-10x for AdamW) and regularization parameters. Consider increasing the `beta` value for very long training runs (e.g., to 0.95 or 0.98).

Warnings

• gotcha Schedule-Free optimizers require explicit calls to `optimizer.train()` and `optimizer.eval()` to manage internal parameter buffers correctly during training and evaluation phases, respectively. Forgetting these calls can lead to incorrect updates or runtime errors.
  Fix: Ensure `optimizer.train()` is called before the training step and `optimizer.eval()` before any evaluation or checkpoint saving. Alternatively, use `ScheduleFreeClosure` versions if your code supports PyTorch Optimizer step closures, as these do not require explicit `train()`/`eval()` calls.
• breaking In version 1.3, the behavior of weight decay during learning rate warmup was changed to improve stability and consistency with standard `AdamW` in PyTorch.
  Fix: If replicating results from versions prior to 1.3, use `AdamWScheduleFreePaper` which retains the older weight decay implementation.
• gotcha If your model utilizes BatchNorm layers, additional modifications are necessary for `test/val` evaluations to function correctly. This is because batch statistics need to be computed from the `x` sequence, not the `y` sequence.
  Fix: Consult the official documentation or examples for specific guidance on handling BatchNorm layers with Schedule-Free optimizers. Using PreciseBN is also suggested to avoid this issue.
• gotcha Training with Schedule-Free optimizers can be more sensitive to the choice of the `beta` parameter than with standard momentum. While the default `0.9` works for many problems, increasing it to `0.95` or `0.98` might be necessary for very long training runs to achieve optimal performance.
  Fix: Experiment with `beta` values like `0.95` or `0.98`, especially for extended training sessions, if initial results are suboptimal.
• gotcha The optimal learning rates for Schedule-Free optimizers are typically higher than those used with schedule-based approaches. For SGD, a learning rate 10x-50x larger might be a good starting point, while for AdamW, 1x-10x larger rates are often effective.
  Fix: Begin hyperparameter tuning with higher learning rates for Schedule-Free optimizers compared to what you would use for their traditional counterparts.

Install

• pip install schedulefree PyPI

Imports

• SGDScheduleFree

  from schedulefree import SGDScheduleFree

• AdamWScheduleFree

  from schedulefree import AdamWScheduleFree

• RAdamScheduleFree

  from schedulefree import RAdamScheduleFree

• ScheduleFreeWrapper

  from schedulefree import ScheduleFreeWrapper

  For wrapping existing PyTorch optimizers to make them schedule-free.

Quickstart

This quickstart demonstrates how to integrate `AdamWScheduleFree` into a basic PyTorch training loop. Key steps include initializing the optimizer, performing forward/backward passes, and crucially, calling `optimizer.train()` and `optimizer.eval()` alongside `model.train()` and `model.eval()` for correct parameter buffer handling during training and evaluation/checkpointing.

import torch
import torch.nn as nn
from schedulefree import AdamWScheduleFree

# 1. Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

# 2. Define a Schedule-Free optimizer (e.g., AdamWScheduleFree)
# Note: Schedule-Free optimizers often benefit from higher learning rates than traditional ones.
optimizer = AdamWScheduleFree(model.parameters(), lr=1e-3, warmup_steps=100)

# 3. Define a loss function
criterion = nn.MSELoss()

# 4. Dummy data for demonstration
inputs = torch.randn(64, 10)
targets = torch.randn(64, 1)

# 5. Training loop (simplified)
num_epochs = 10
for epoch in range(num_epochs):
    model.train() # Standard PyTorch model training mode
    optimizer.train() # REQUIRED for Schedule-Free optimizers

    # Forward pass
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 2 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

    # Evaluation phase (e.g., for validation or checkpointing)
    model.eval() # Standard PyTorch model evaluation mode
    optimizer.eval() # REQUIRED for Schedule-Free optimizers before evaluation/checkpointing
    with torch.no_grad():
        val_inputs = torch.randn(16, 10)
        val_targets = torch.randn(16, 1)
        val_outputs = model(val_inputs)
        val_loss = criterion(val_outputs, val_targets)
        # In a real scenario, you'd calculate metrics on val_outputs and val_targets

print("Training complete.")