Struct AdamW 

pub struct AdamW {
    pub lr: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub eps: f32,
    pub weight_decay: f32,
    pub f32_math: bool,
    /* private fields */
}

Adam with decoupled weight decay.

Per-tensor state identical to crate::Adam (two f32 buffers).

Fields

lr: f32

Learning rate. Typical LLM pre-training value: 1e-4 to 3e-4.

beta1: f32

First-moment EMA decay. Default 0.9.

beta2: f32

Second-moment EMA decay. Default 0.999 (matches Adam); 0.95 is common for very long pre-training runs.

eps: f32

Denominator stability constant. Default 1e-8.

weight_decay: f32

Decoupled weight-decay coefficient λ. Multiplies the parameter directly inside the update; 0.01–0.1 typical. Defaults to 0.01.

f32_math: bool

When true, moment/parameter updates stay in pure f32 (PyTorch AdamW / foreach=False default). Default false uses f64 intermediates (slightly more accurate, but drifts vs PyTorch shared-init checks).

Implementations

impl AdamW

pub fn new(lr: f32) -> Self

Construct with the given learning rate and the standard (β₁, β₂, ε, λ) = (0.9, 0.999, 1e-8, 0.01) defaults.

Examples found in repository

examples/step_bench.rs (line 63)

fn main() {
    for n in [64 * 1024usize, 1024 * 1024, 8 * 1024 * 1024] {
        println!(
            "\n{} elements ({:.1} MiB per buffer)",
            n,
            (n * 4) as f64 / 1048576.0
        );
        bench("adamw (f64 default)", n, Box::new(AdamW::new(1e-3)));
        bench(
            "adamw f32Math",
            n,
            Box::new(AdamW::new(1e-3).with_f32_math(true)),
        );
        bench("adam (f64 default)", n, Box::new(Adam::new(1e-3)));
        bench(
            "adam f32Math",
            n,
            Box::new(Adam::new(1e-3).with_f32_math(true)),
        );
        bench("sgd+momentum", n, {
            let mut o = Sgd::new(1e-3);
            o.momentum = 0.9;
            Box::new(o)
        });
        bench("lion", n, Box::new(Lion::new(1e-3)));
    }
}

pub fn with_betas(self, b1: f32, b2: f32) -> Self

Override (β₁, β₂).

pub fn with_weight_decay(self, wd: f32) -> Self

Override the decoupled-decay coefficient.

pub fn with_eps(self, eps: f32) -> Self

Override the denominator ε.

pub fn with_f32_math(self, on: bool) -> Self

Pure-f32 AdamW arithmetic (match PyTorch training trajectories).

Examples found in repository

examples/step_bench.rs (line 67)

fn main() {
    for n in [64 * 1024usize, 1024 * 1024, 8 * 1024 * 1024] {
        println!(
            "\n{} elements ({:.1} MiB per buffer)",
            n,
            (n * 4) as f64 / 1048576.0
        );
        bench("adamw (f64 default)", n, Box::new(AdamW::new(1e-3)));
        bench(
            "adamw f32Math",
            n,
            Box::new(AdamW::new(1e-3).with_f32_math(true)),
        );
        bench("adam (f64 default)", n, Box::new(Adam::new(1e-3)));
        bench(
            "adam f32Math",
            n,
            Box::new(Adam::new(1e-3).with_f32_math(true)),
        );
        bench("sgd+momentum", n, {
            let mut o = Sgd::new(1e-3);
            o.momentum = 0.9;
            Box::new(o)
        });
        bench("lion", n, Box::new(Lion::new(1e-3)));
    }
}

Trait Implementations

impl Clone for AdamW

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for AdamW

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Optimizer for AdamW

fn set_lr(&mut self, lr: f32)

Set the base learning rate (for LR schedules / warmup). Default is a no-op for algorithms without a scalar lr (e.g. Adafactor’s relative step size); every algorithm in this crate that has an lr field overrides this to update it.

fn state_dict(&self) -> Option<OptimizerState>

Snapshot the optimizer’s state for a checkpoint.

fn load_state_dict(&mut self, state: &OptimizerState) -> bool

Restore a snapshot. Returns false when unsupported or when the state does not belong to this algorithm.

fn step( &mut self, name: &str, _shape: &[usize], param: &mut [f32], grad: &[f32], )

fn end_iteration(&mut self)

Advance the global step counter. Most algorithms increment per call to step, so most implementations leave this a no-op.

fn step_batch(&mut self, items: &mut [OptItem<'_>])

Batched step over ALL parameters in one call. Default: sequential step per item — bit-identical to the per-parameter loop. Optimizers whose parameter groups are independent (e.g. Muon on the 2-D weight matrices vs AdamW on the embeddings/biases/norms) can override this to run the groups on separate threads; because the groups touch disjoint parameters and disjoint optimizer state, the result is bit-for-bit the same as the serial loop — only the wall-clock (the CPU-side optimizer bubble) shrinks toward max(group_times) instead of their sum.

fn lr_scale(&self, _name: &str) -> f32

Per-tensor multiplier on the effective learning rate. Default is 1.0 for every name. Override when wrapping this crate to support per-name LR schedules (e.g. embedding-vs-attention splits, or the Gaussian-splat attribute-typed LR setup). The CPU impls in this crate currently honor this only when the caller passes a pre-scaled lr for the relevant call — backends are encouraged to consult it inside their fused kernel.