getting-started.md

Getting Started

Installation

Use Python 3.10 or later with PyTorch 2.1 or later. The package installer selects dependencies that support your Python version.

pip install pytorch-optimizer

Install optional integrations from their projects as needed: bitsandbytes, Q-GaLore, or TorchAO.

Run a training step

Pass the model parameters to an optimizer class. Use the optimizer in a PyTorch training loop:

import torch
from torch import nn
from pytorch_optimizer import AdamP

model = nn.Linear(10, 1)
optimizer = AdamP(model.parameters(), lr=1e-3)
inputs = torch.randn(8, 10)
targets = torch.randn(8, 1)

optimizer.zero_grad()
loss = nn.functional.mse_loss(model(inputs), targets)
loss.backward()
optimizer.step()

Load an optimizer by name

Use load_optimizer() to select an optimizer by its name in a training configuration:

from pytorch_optimizer import load_optimizer

optimizer = load_optimizer('adamp')(model.parameters(), lr=1e-3)

Use Torch Hub to load an optimizer class:

optimizer_class = torch.hub.load('kozistr/pytorch_optimizer', 'adamp')
optimizer = optimizer_class(model.parameters(), lr=1e-3)

Combine optimizer options

Pass the model to create_optimizer() to configure weight decay and optimizer wrappers:

from pytorch_optimizer import create_optimizer

optimizer = create_optimizer(
    model,
    'adamp',
    lr=1e-3,
    weight_decay=1e-3,
    use_gc=True,
    use_lookahead=True,
)

Check the optimizer reference for options and arguments specific to each optimizer.

Compile optimizer steps

Set compile=True to let PyTorch optimize and automatically fuse optimizer updates on CPU or GPU.

optimizer = create_optimizer(model, 'laprop', lr=1e-3, compile=True)

APOLLO, Conda, Fira, and GaLore support this option for unprojected parameters. FlashAdamW supports it for unquantized FP32 parameters with decouple_lr=False. EmoNavi, EmoLynx, and EmoFact keep FP16 and BF16 updates eager. SPlus compiles nonmatrix updates and keeps matrix preconditioning eager.

The factory selects the batched update path automatically for supported optimizers. Matrix updates can batch compatible flattened shapes, including convolution weights and tall/wide pairs, while preserving learning-rate scaling. The compiler chooses which operations to fuse; no separate fusion setting is required. State initialization, step counters, and scalar bookkeeping stay eager on these paths. The update receives tensor learning rates on each parameter device, so changing the rate does not require recompilation. Groups with unsupported options or complex parameters use eager per-parameter updates. Other optimizers and optimizer wrappers compile the whole step. Graph breaks and recompilation can vary by optimizer.

Use compile=False (the default) for eager execution. Use eager execution on Python 3.15 because PyTorch disables compilation for that version. Pass torch.compile() options through compile_kwargs. Run with TORCH_LOGS=graph_breaks,recompiles to inspect graph breaks and recompilation.

Advanced execution controls

Most users only need compile=True. Existing optimizer-specific settings remain available for benchmarking and workloads that need an explicit backend.

These options select different execution strategies for the same optimizer algorithm:

Option	Behavior
foreach=True	Batch each operation across multiple parameter tensors.
compile=True	Compile tensor operations, allowing the compiler to fuse updates into fewer kernels.
fused=True	Select an optimizer's dedicated fused implementation, when available.

For dedicated compiled foreach updates, the factory enables foreach=True unless you set foreach=False. With foreach=False, it compiles the whole step and converts a float learning rate to a tensor on the model's device. For native AdamW on CUDA, set capturable=True if you use a tensor rate with foreach=True.

Fusion combines operations such as momentum accumulation, normalization, weight decay, and parameter updates. It reduces kernel launches and intermediate memory traffic. Fused kernels can be written in C++/CUDA or Triton, or generated by a compiler; the term describes how operations execute, rather than the implementation language. See PyTorch's optimizer implementation notes and optimizer compilation tutorial.

The factory forwards optimizer-specific options, so native AdamW can use its existing fused implementation:

optimizer = create_optimizer(model, 'adamw', lr=1e-3, fused=True)

Check the optimizer and PyTorch version for supported devices and dtypes. Compilation alone does not require a fused argument. Benchmark after compilation warmup to compare performance.

Train Muon with FSDP2

FSDP2 splits model parameters, gradients, and optimizer states across GPUs to save memory. Use CUDA-enabled PyTorch with FSDP2 support. Each process uses one GPU.

In your training script, initialize distributed training, then create the optimizer after sharding the model:

import os

import torch
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard

from pytorch_optimizer import create_optimizer

local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
dist.init_process_group('nccl')

model = YourModel().cuda(local_rank)
fully_shard(model)
optimizer = create_optimizer(model, 'muon', lr=0.02, foreach=True)

for inputs, targets in train_loader:
    optimizer.zero_grad()
    loss = loss_fn(model(inputs.cuda(local_rank)), targets.cuda(local_rank))
    loss.backward()
    optimizer.step()
dist.destroy_process_group()

Use your usual model and loss function. Build train_loader with DistributedSampler after initializing the process group so each rank receives different data. For large models, call fully_shard() on each transformer block before sharding the root model. Keep optimizer steps eager (compile=False, the default).

Run on one node with two GPUs:

torchrun --standalone --nproc-per-node=2 train.py

For two nodes with two GPUs each, run the same command on both nodes. Replace HOST with a hostname or IP reachable from both nodes:

torchrun --nnodes=2 --nproc-per-node=2 --rdzv-backend=c10d --rdzv-endpoint=HOST:29400 --rdzv-id=muon train.py

For AdaMuon, use 'adamuon' in create_optimizer() with the same setup.

MuonClip also uses the same setup. Label its query/key groups and reduce per-head attention maxima over all accumulation microbatches and data-parallel ranks with ReduceOp.MAX before calling step(attention_max_logits=...).

Discover components

Filter component names with wildcard patterns:

from pytorch_optimizer import (
    get_supported_loss_functions,
    get_supported_lr_schedulers,
    get_supported_optimizers,
)

all_optimizers = get_supported_optimizers()
adam_family = get_supported_optimizers('adam*')
selected = get_supported_optimizers(['adam*', 'ranger*'])
cosine_schedulers = get_supported_lr_schedulers('cosine*')
focal_losses = get_supported_loss_functions('*focal*')

Check the scheduler reference and loss reference for signatures and usage details. Read the FAQ before you train with Hessian-based optimizers such as AdaHessian and SophiaH.