Optimizers

Configuring optimizers

Overview

Axolotl supports all optimizers supported by transformers OptimizerNames

Here is a list of optimizers supported by transformers as of v4.54.0:

• adamw_torch
• adamw_torch_fused
• adamw_torch_xla
• adamw_torch_npu_fused
• adamw_apex_fused
• adafactor
• adamw_anyprecision
• adamw_torch_4bit
• adamw_torch_8bit
• ademamix
• sgd
• adagrad
• adamw_bnb_8bit
• adamw_8bit # alias for adamw_bnb_8bit
• ademamix_8bit
• lion_8bit
• lion_32bit
• paged_adamw_32bit
• paged_adamw_8bit
• paged_ademamix_32bit
• paged_ademamix_8bit
• paged_lion_32bit
• paged_lion_8bit
• rmsprop
• rmsprop_bnb
• rmsprop_bnb_8bit
• rmsprop_bnb_32bit
• galore_adamw
• galore_adamw_8bit
• galore_adafactor
• galore_adamw_layerwise
• galore_adamw_8bit_layerwise
• galore_adafactor_layerwise
• lomo
• adalomo
• grokadamw
• schedule_free_radam
• schedule_free_adamw
• schedule_free_sgd
• apollo_adamw
• apollo_adamw_layerwise
• stable_adamw

Custom Optimizers

Enable custom optimizers by passing a string to the optimizer argument. Each optimizer will receive beta and epsilon args, however, some may accept additional args which are detailed below.

optimi_adamw

optimizer: optimi_adamw

ao_adamw_4bit

Deprecated: Please use adamw_torch_4bit.

ao_adamw_8bit

Deprecated: Please use adamw_torch_8bit.

ao_adamw_fp8

optimizer: ao_adamw_fp8

adopt_adamw

GitHub: https://github.com/iShohei220/adopt Paper: https://arxiv.org/abs/2411.02853

optimizer: adopt_adamw

adamc

AdamW with the weight decay corrected for the learning rate schedule. AdamW’s decoupled decay makes the steady-state gradient-to-weight-norm ratio depend on the current learning rate, so a decaying schedule pushes the gradient norm up over the second half of training. AdamC scales the decay by the schedule position (lambda_t = lambda * lr_t / lr_max), which holds that ratio constant and keeps weight norms stable.

Paper: https://arxiv.org/abs/2506.02285

optimizer: adamc
learning_rate: 0.00001
weight_decay: 0.05
lr_scheduler: cosine   # the correction is a no-op under a constant LR

The peak learning rate is captured per parameter group when the optimizer is built, so embedding_lr, embedding_lr_scale, and lr_groups are each corrected against their own peak. Override it explicitly if you build a schedule that exceeds the configured LR:

optim_args:
  max_lr: 0.0001

Only groups with a non-zero weight decay are affected — axolotl already excludes embeddings, the LM head, norms, and biases from decay, matching the paper’s rule of applying the correction to normalized layers only.

came_pytorch

GitHub: https://github.com/yangluo7/CAME/tree/master Paper: https://arxiv.org/abs/2307.02047

optimizer: came_pytorch

# optional args (defaults below)
adam_beta1: 0.9
adam_beta2: 0.999
adam_beta3: 0.9999
adam_epsilon: 1e-30
adam_epsilon2: 1e-16

muon

Blog: https://kellerjordan.github.io/posts/muon/ Paper: https://arxiv.org/abs/2502.16982v1

optimizer: muon

dion

Microsoft’s Dion (DIstributed OrthoNormalization) optimizer is a scalable and communication-efficient orthonormalizing optimizer that uses low-rank approximations to reduce gradient communication.

GitHub: https://github.com/microsoft/dion Paper: https://arxiv.org/pdf/2504.05295 Note: Implementation written for PyTorch 2.7+ for DTensor

optimizer: dion
dion_lr: 0.01
dion_momentum: 0.95
lr: 0.00001  # learning rate for embeddings and parameters that fallback to AdamW

sinkgd

SinkGD (Gradient Multi-Normalization) is a stateless optimizer: 2D linear weight matrices are updated via the SR-Sinkhorn procedure (alternating row/column L2 normalization of the raw gradient, no momentum or variance state), while embeddings, the LM head, and 1D params (norms/biases) fall back to AdamW. The AdamW fallback uses torchao’s 8-bit optimizer base, so optimizer-state memory is very small (~87% less than 8-bit AdamW on an 8B full finetune, since the ~87% of params that are 2D linear carry zero optimizer state).

Requires PyTorch >= 2.5.1 (relies on torchao’s low-bit optimizer base and torch.compile).

Paper: https://arxiv.org/abs/2502.06742

optimizer: sinkgd
learning_rate: 0.001
optim_args:
  sinkhorn_iters: 5      # number of SR-Sinkhorn row/column normalization iterations
  sinkgd_lr_scale: 0.05  # α scale applied to linear-layer updates

Width transfer (sinkgd_base_width) and spectral normalization

Two optional, default-off knobs help the learning rate transfer across model width and condition the update. sinkgd_base_width and sinkgd_spectral_target: muon are mutually exclusive — both correct for width, so enabling both double-counts (validated and rejected at construction). Spectral norm at the default unit target combines fine with sinkgd_base_width, since it conditions the update without touching width scaling.

Width-aware scaling (sinkgd_base_width). SinkGD’s update is Adam-class: its per-layer learning rate should scale as 1/d_in. Set sinkgd_base_width to the hidden size you tuned sinkgd_lr_scale on, and each 2D-linear update is scaled by alpha_eff = sinkgd_lr_scale * (base_width / d_in) ** sinkgd_lr_width_exponent (d_in is the layer’s input dim — the shared input of a fused QKV / gate-up matrix). With sinkgd_base_width unset, behavior is identical to before (plain scalar sinkgd_lr_scale), so existing configs are unchanged.

optim_args:
  sinkgd_lr_scale: 0.05
  sinkgd_base_width: 2048        # d_in you tuned sinkgd_lr_scale at; enables 1/d_in transfer
  sinkgd_lr_width_exponent: 1.0  # 1.0 = pure 1/d_in (default); tune only if width varies >10x

Spectral normalization (sinkgd_spectral_norm). Rescales each update to a target operator norm via a cheap warm-started power iteration (~2 matvecs/layer, +1–3% step time, one O(d_in) state vector per matrix). With sinkgd_spectral_target: muon it pins the operator norm to sqrt(d_out/d_in), which makes SinkGD Muon-class and both transfers across width and trained to lower loss than plain 1/d_in in width sweeps — so use it instead of sinkgd_base_width, not with it. unit (the default target) pins a width-independent norm and acts as a pure conditioning stabilizer on top of 1/d_in.

optim_args:
  sinkgd_lr_scale: 0.05
  sinkgd_spectral_norm: true
  sinkgd_spectral_target: muon   # spectral norm owns width transfer (leave sinkgd_base_width unset)
  sinkgd_spectral_norm_iters: 1  # power-iteration steps per update (warm-started; 1–2 is enough)

Distributed (fsdp_version: 2): width-aware scaling and spectral norm both work under FSDP2/TP sharding. On replicated / expert-sharded weights everything runs locally with no extra communication. On a matrix-dim-sharded weight (the common FSDP2 row-sharded case) the spectral norm’s power iteration adds one small vector all-reduce (length d_in or d_out) per iteration over the dp_shard group — the same shape and group as SR-Sinkhorn’s existing norm-vector reduce, and the matrix itself is never gathered. The persisted power-iteration vector round-trips through FSDP2 checkpoints.

Fused Triton kernels (sinkgd_fused_kernel)

sinkgd_fused_kernel: true replaces the compiled update with fused Triton kernels (one kernel per SR-Sinkhorn iteration; the column scale, power iteration, and weight update fold into the same passes). Works for all SinkGD variants (plain, spectral norm, MD sphere) on single-device and FSDP2 rows-sharded weights — same all-reduce count as the compiled path, with an automatic tall/wide grid layout so both full matrices and heavily-sharded wide-short local shards stay fully occupied. Measured on B200: 1.4–1.75x on the optimizer step single-GPU, 1.0–1.8x on 2-rank FSDP2, all regimes ≥1x with the wide layout. Numerics are equivalent at bf16 rounding scale but not byte-identical to the compiled path, so the flag is off by default. Falls back to the compiled path for cols-sharded (TP) weights, bf16_stochastic_round, or when Triton is unavailable. Spectral-norm / MD-sphere updates that run without a shard group to amortize against (single device, or replicated / expert-sharded weights under FSDP2) also fall back when the matrix has fewer than sinkgd_fused_min_numel elements (default 2^25) — at those sizes the epilogue’s extra kernel launches make the compiled path faster.

optim_args:
  sinkgd_fused_kernel: true

Weight-sphere variant (sinkgd_md_sphere, experimental)

sinkgd_md_sphere: true switches to an experimental magnitude–direction variant: each SinkGD-routed 2D weight is held on a fixed Frobenius sphere (||W||_F anchored at enable time), the SR-Sinkhorn update is spectral-normalized to unit operator norm, applied, and the weight is reprojected onto the sphere (no learnable gains — SinkGD’s row/column balancing makes them redundant). In width/depth sweeps this bounds the deepest-layer activation growth most tightly of any variant, but its optimal LR is width-dependent (lr_opt ∝ d_model**-0.6) and its optimum is narrow, so it needs per-scale LR tuning and is off by default. Prefer the plain spectral-norm path (above) unless you specifically want the tightest activation bound. It adds a per-matrix scalar all-reduce (the sphere’s global Frobenius norm) on top of the spectral all-reduce under FSDP2; the sphere radius and power-iteration vector round-trip.

optim_args:
  sinkgd_md_sphere: true
  sinkgd_lr_scale: 0.05
  # tune learning_rate per width: lr_opt ~ d_model**-0.6 (e.g. ~0.03 at d=2048)

q_galore_adamw8bit

Q-GaLore extends GaLore with two extra ideas: an INT4-quantized projection matrix and an adaptive SVD scheduler that skips re-projection when a layer’s gradient subspace stabilizes. Both are wired up in axolotl. The third Q-GaLore trick — INT8 weight wrapping — is not yet implemented and is tracked as a follow-up.

GitHub: https://github.com/VITA-Group/Q-GaLore Paper: https://arxiv.org/abs/2407.08296

Install: pip install axolotl[qgalore]

This optimizer is for full fine-tuning. It is incompatible with adapter (LoRA/QLoRA), load_in_8bit, and load_in_4bit. DeepSpeed is currently gated off; FSDP requires fsdp_version: 2 with use_orig_params: true.

optimizer: q_galore_adamw8bit
bf16: true

# which parameter substrings get the low-rank projection
# (defaults to ["attn", "mlp"] if unset — matches the reference impl)
optim_target_modules:
  - attn
  - mlp

# Q-GaLore hyperparameters (defaults shown)
qgalore_rank: 256
qgalore_update_proj_gap: 200       # max steps between SVD refreshes
qgalore_scale: 0.25
qgalore_proj_type: std
qgalore_proj_quant: true           # INT-quantize the projection matrix P
qgalore_proj_bits: 4               # bitwidth for P
qgalore_proj_group_size: 256       # must divide P's last dim evenly
qgalore_cos_threshold: 0.4         # skip SVD if P_t is this similar to P_{t-1}
qgalore_gamma_proj: 2              # grow update_proj_gap by this factor when stable
qgalore_queue_size: 5

polora

PoLoRA (Preconditioned Orthogonalized LoRA) is a LoRA-only optimizer. It extends Muon-style spectral descent to the (A, B) product structure: each pair is whitened in a diagonal-Kronecker curvature metric, passed through the matrix sign, then rescaled to spectral norm rho = lr / (sigma_max(A) + sigma_max(B)). Optimizer state is a first moment plus two diagonal preconditioners, roughly half of Adam’s.

GitHub: https://github.com/nikhilgsh/polora Docs: https://nikhilgsh.github.io/polora/

Install from source:

pip install git+https://github.com/nikhilgsh/polora.git

• Requires adapter: lora or adapter: qlora, and single-GPU, DDP, or FSDP2.
• We recommend enabling sample_packing and higher batch sizes to minimize optim step overhead.
• We find that learning_rate should be set 1-2 orders of magnitude higher than AdamW: 3e-4 -> 3e-2.

optimizer: polora
adapter: lora
learning_rate: 0.01      # make sure to sweep

# optional, defaults shown
optim_args:
  beta1: 0.9             # momentum
  curvature_beta: 0.99   # EMA for the diagonal preconditioners
  epsilon: 1.0e-12       # numerical floor
  delta: 1.0e-4          # relative damping of the inverse square roots
  ns_steps: 8            # PolarExpress matrix-sign iterations
  higham_iters: 8        # Newton-Schulz iterations for the inverse square roots
  compile: false         # torch.compile the spectral kernels

Note

Memory. PoLoRA carries one momentum where Adam carries two, but stacks same-shape pairs into fp32 temporaries peaking near 6x the size of the LoRA factors. At 1B those transients fit under the activation peak, and it came out slightly cheaper than AdamW at both lora_r: 32 and 256. They scale with lora_r times model width, so on a wide model at high rank will be more costly.

Gradients. Every LoRA factor needs a gradient every step. Exclude layers that skip batches (vision towers, unrouted MoE experts) via lora_exclude_modules. On MoE, lora_target_parameters builds factors at lora_r * num_experts, so keep lora_r small.

FSDP2. Axolotl adds FSDP2 support, however, we recommend only using it when frozen base does not fit due to noticeable slowdowns.