--- title: Optimizers description: Configuring optimizers --- ## Overview Axolotl supports all optimizers supported by [transformers OptimizerNames](https://github.com/huggingface/transformers/blob/51f94ea06d19a6308c61bbb4dc97c40aabd12bad/src/transformers/training_args.py#L142-L187) 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 ```yaml optimizer: optimi_adamw ``` ### ao_adamw_4bit Deprecated: Please use `adamw_torch_4bit`. ### ao_adamw_8bit Deprecated: Please use `adamw_torch_8bit`. ### ao_adamw_fp8 ```yaml optimizer: ao_adamw_fp8 ``` ### adopt_adamw GitHub: [https://github.com/iShohei220/adopt](https://github.com/iShohei220/adopt) Paper: [https://arxiv.org/abs/2411.02853](https://arxiv.org/abs/2411.02853) ```yaml optimizer: adopt_adamw ``` ### came_pytorch GitHub: [https://github.com/yangluo7/CAME/tree/master](https://github.com/yangluo7/CAME/tree/master) Paper: [https://arxiv.org/abs/2307.02047](https://arxiv.org/abs/2307.02047) ```yaml 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/](https://kellerjordan.github.io/posts/muon/) Paper: [https://arxiv.org/abs/2502.16982v1](https://arxiv.org/abs/2502.16982v1) ```yaml 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](https://github.com/microsoft/dion) Paper: [https://arxiv.org/pdf/2504.05295](https://arxiv.org/pdf/2504.05295) Note: Implementation written for PyTorch 2.7+ for DTensor ```yaml 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](https://arxiv.org/abs/2502.06742) ```yaml 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. ```yaml 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`. ```yaml 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. ```yaml 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. ```yaml 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](https://arxiv.org/abs/2403.03507) 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](https://github.com/VITA-Group/Q-GaLore) Paper: [https://arxiv.org/abs/2407.08296](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`. ```yaml 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 ```