torch.optim

Active contributors: albanD, janeyx99

Purpose

torch.optim provides numerical optimizers (SGD, Adam, AdamW, RMSprop, Adafactor, LBFGS, …) and learning-rate schedulers. Each optimizer is a class that owns parameter groups, hyperparameters, and per-step state (momentum buffers, running averages, …).

Directory layout

Path	Contents
torch/optim/__init__.py	Public surface
torch/optim/optimizer.py	The Optimizer base class
torch/optim/sgd.py	SGD
torch/optim/adam.py	Adam, AdamW
torch/optim/adamw.py	AdamW (the proper decoupled-weight-decay variant)
torch/optim/rmsprop.py	RMSprop
torch/optim/lbfgs.py	L-BFGS
torch/optim/adafactor.py	Adafactor
torch/optim/lr_scheduler.py	LR schedulers
torch/optim/swa_utils.py	Stochastic Weight Averaging utilities
torch/distributed/optim/	Distributed optimizers (ZeroRedundancyOptimizer, OSS, FSDP-friendly)

Each optimizer file defines two functions:

• The public class (e.g., Adam) with step() taking an optional closure.
• A functional _single_tensor_* and _multi_tensor_* (and sometimes _fused_*) implementation. The class delegates to one of them based on availability.

Key abstractions

Type	File	Purpose
Optimizer	torch/optim/optimizer.py	Base class
Parameter group	dict in optimizer.param_groups	Per-group hyperparameters
Optimizer state	optimizer.state[param]	Per-parameter persistent state (momentum, …)
_LRScheduler	torch/optim/lr_scheduler.py	Base class for schedulers

How it works

Step

A typical step:

optimizer.zero_grad()
loss.backward()
optimizer.step()

Inside step:

1. Iterate over parameter groups.
2. For each parameter with a .grad, fetch its state.
3. Apply the update rule (multi-tensor or fused, where available).
4. Optionally invoke a closure for line-search optimizers.

Implementations: single, multi, fused

Each optimizer ships up to three implementations:

• _single_tensor_* — vanilla Python loop over parameters. Easy to read, slow.
• _multi_tensor_* — uses torch._foreach_* ops to update many parameters per kernel call. The default for CUDA.
• _fused_* — single-kernel CUDA implementation (Adam, AdamW). The fastest.

The choice is controlled by the foreach= and fused= constructor arguments, with sensible defaults.

State dict

optimizer.state_dict() and optimizer.load_state_dict(...) serialize the per-parameter state and the param_groups. Distributed checkpoint handles the sharded version.

Capturable / differentiable optimizers

For research use cases (meta-learning, MAML, hypernetwork training), most optimizers support differentiable=True (treats the step as a differentiable function) and capturable=True (lets state live on GPU so the step can be CUDA-graphed).

LR schedulers

torch/optim/lr_scheduler.py ships:

• StepLR, MultiStepLR, ExponentialLR, CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau, LinearLR, CyclicLR, LambdaLR, SequentialLR, ChainedScheduler, …

Each subclasses _LRScheduler. scheduler.step() is called once per epoch (or per batch for some) and updates optimizer.param_groups[i]['lr'].

SWA utilities

torch/optim/swa_utils.py implements Stochastic Weight Averaging (SWA) and Exponential Moving Average (EMA) wrappers. Used commonly for ensemble-like inference.

Distributed

torch/distributed/optim/ ships ZeRO-style optimizers that shard optimizer state across ranks (ZeroRedundancyOptimizer, OSS). Most modern jobs use FSDP, which subsumes this; the legacy classes remain for backwards compatibility.

Where to look

File	Purpose
torch/optim/optimizer.py	Base class
torch/optim/adam.py	Most-studied optimizer
torch/optim/adamw.py	The default choice for transformers
torch/optim/lr_scheduler.py	Schedulers
torch/optim/swa_utils.py	SWA / EMA

Where to read next

• Packages / torch.nn — modules whose parameters get optimized.
• Features / Distributed training — sharded optimizers.