dew.training.optim

The optimizer a recipe builds from an OptimConfig.

Every recipe wires the same solver: a warmup-cosine schedule when one is asked for, weight decay folded into the optimizer’s own kwargs, and global-norm clipping. That wiring is library behavior, so it lives here and the recipes call it. The Trainer forms the normalized effective-window gradient before calling this solver.

The ‘muon’ entry is the production parameter-group split the labs converged on (docs/research/frontier-training.md:183). AdamW takes the embeddings, the head, the router and the norms; Muon takes the matrices.

optax.contrib.muon owns the masked composition, partitioning with optax.masked per group (optax/contrib/_muon.py:694). What Dew supplies is the parameter spec: which group a parameter belongs to, and which of its axes are the matrix.

Name	Summary
HEAD_AXES
BATCH_AXES
SELECTION_AXES
muon_weight_dimension_numbers	Build a MuonDimensionNumbers per parameter, None where AdamW steps in.
QK_PROJECTIONS	The projection names the clip rescales: the query and key projections of grouped-query attention and, for latent attention, the per-head output projections.
scale_by_qk_clip	Rescale the query and key projections of every head past tau.
stochastic_round_bf16	x rounded to bf16 up or down with the probability of its distance to each, from a hash of seed and each element’s flat index.
bf16_moments	optax.scale_by_adam with both moments stored in bf16.
BF16_STATE_OPTIMIZERS
OPTIMIZER_MAP
ParamGroup	Parameters the optimizer moves at their own learning rate and decay.
NO_DECAY_PATTERNS	What lm-engine’s no_weight_decay group holds (configs/param-groups/ mup.yml at 45b6b57b): biases, every norm’s weight (an RMSNorm keeps scale here, Mamba-2’s gated norm weight) and Mamba-2’s dt_bias.
mup_param_groups	lm-engine’s muP parameter groups (configs/param-groups/mup.yml at 45b6b57b), in its order: norms, biases and dt_bias without weight decay at the base rate; the token embeddings at the base rate with decay; everything else, the router, A_log, D and the conv taps included, at the base rate divided by width_multiplier (lm-engine’s m_width, the model’s logits_scaling).
param_labels	The optax.multi_transform labeller: each leaf’s first matching group.
power_schedule	lm-engine’s power scheduler (optimization/lr_scheduler/power.py at 45b6b57b) with an optional linear tail.
linear_schedule	lm-engine’s linear scheduler (lr_scheduler/linear.py at 45b6b57b): from zero to peak over the warmup, constant to decay_start (None: the warmup’s end), then linear to end_value at decay_end.
ScheduleBase	One learning-rate schedule’s record: its own fields and its optax schedule.
Cosine	Linear warmup from init to peak, cosine to end at decay_steps (None: the run’s end); optax.warmup_cosine_decay_schedule.
PowerTail	A power schedule’s linear tail: from the law’s rate at start to end at step steps (None: the run’s end).
Power	lm-engine’s power scheduler, power_schedule: warmup, then min(peak, a * (step * c) ** b), and with a tail a linear decay after the law (Rigel’s last 29%).
Linear	lm-engine’s linear scheduler, linear_schedule: warmup to peak, constant to decay_start (None: the warmup’s end), linear to end at decay_steps (None: the run’s end).
learning_rate_schedule	The rate config names: its schedule over a steps-update run, or the constant learning_rate when it names none.
build_optimizer	Build the solver a config describes, with its schedule, parameter groups and clipping.

HEAD_AXES

attribute source

HEAD_AXES = frozenset({'heads', 'head_dim', 'kv'})

BATCH_AXES

attribute source

BATCH_AXES = frozenset({'exp'})

SELECTION_AXES

attribute source

SELECTION_AXES = frozenset({'vocab', 'output'})

muon_weight_dimension_numbers

function source

def muon_weight_dimension_numbers(params)

Build a MuonDimensionNumbers per parameter, None where AdamW steps in.

Which group a parameter lands in is read off the logical axes its module declares (dew.nn.sharding), the same table the sharding derivation reads, so one declaration answers both questions.

AdamW takes a parameter of rank below two, a bias, and a parameter that maps into or out of a discrete index: the vocabulary, the model’s output space, or the expert a router picks. That is the split four labs cross-confirmed. Everything else is a matrix and goes to Muon.

An undeclared matrix of rank two takes Linen’s kernel convention, contracting axis 0 into axis 1. An undeclared parameter of higher rank raises, because its matrix axes are what this spec cannot guess, and orthogonalizing the wrong pair would show up as a worse loss curve.

Optax reads one spec tree shaped like the parameters and treats a None leaf as an AdamW parameter (optax/contrib/_muon.py:660-675).

QK_PROJECTIONS

attribute source

QK_PROJECTIONS = frozenset({'q_proj', 'q_b_proj', 'k_proj', 'kv_b_proj'})

The projection names the clip rescales: the query and key projections of grouped-query attention and, for latent attention, the per-head output projections. Anything else keeps its update untouched.

scale_by_qk_clip

function source

def scale_by_qk_clip(tau: float = 100.0) -> optax.GradientTransformationExtraArgs

Rescale the query and key projections of every head past tau.

This is Kimi K2’s MuonClip (arXiv 2507.20534), applied after the update.

The per-head maxima arrive as qk_stats, the qk collection the model sowed, which the trainer forwards from the loss’s Aux. Without them the transform steps aside, leaving every other optimizer and every run whose loss never opened the collection on its old update.

stochastic_round_bf16

function source

def stochastic_round_bf16(x: jax.Array, seed: jax.Array) -> jax.Array

x rounded to bf16 up or down with the probability of its distance to each, from a hash of seed and each element’s flat index.

Adding 16 uniform bits below the kept mantissa and truncating rounds up exactly when the discarded bits plus the noise carry, which is the discarded fraction’s probability. The noise is a counter-based hash, so a step’s rounding is a pure function of the seed and the position: no key is split or carried, and nothing is read from memory. jax.random.bits (threefry) made the whole update 1.75x slower than fp32 state on a TPU v6e and 1.05x on an L4; this costs a few integer ops per element. NaN stays NaN.

bf16_moments

function source

def bf16_moments(
    b1: float,
    b2: float,
    eps: float,
    eps_root: float,
    nesterov: bool,
) -> optax.GradientTransformation

optax.scale_by_adam with both moments stored in bf16.

Each update runs optax’s own update one leaf at a time, on that leaf’s moments widened to fp32, and writes the new moments back stochastically rounded (stochastic_round_bf16), seeded by the step count, the leaf and the moment: the step is optax’s, only the storage is Dew’s. Leaf by leaf keeps one leaf’s fp32 moments live at a time; widening the whole state first held all of them and raised the update’s peak from 4.49 to 7.37 GB on the lm-dense tree (RTX 4080). Round to nearest would lose every increment of the second moment smaller than half its bf16 spacing, which at b2 = 0.999 is most of them; a stochastic rounding keeps each one in expectation. The state keeps optax’s ScaleByAdamState layout, so sharding and checkpoints read it as they read fp32 state.

BF16_STATE_OPTIMIZERS

attribute source

BF16_STATE_OPTIMIZERS = {'adam': _bf16_adam, 'adamw': _bf16_adamw}

OPTIMIZER_MAP

attribute source

OPTIMIZER_MAP = {'adam': optax.adam, 'adamw': optax.adamw, 'lamb': optax.lamb, 'muon': _muon_groups, 'muonclip': _muonclip_groups}

ParamGroup

dataclass source

class ParamGroup(
    name: str,
    patterns: tuple[str, ...],
    learning_rate_multiplier: float = 1.0,
    weight_decay: float | None = None,
)

Parameters the optimizer moves at their own learning rate and decay.

patterns are fnmatch patterns over a parameter’s path, its dict keys joined by ’/’ (layers_3/self_attn/q_proj/kernel); * crosses ’/’. A parameter joins the first group of OptimConfig.param_groups a pattern of which it matches, lm-engine’s rule (optimization/params_group.py at 45b6b57b), and one that matches none raises. The group’s learning rate is the schedule’s times learning_rate_multiplier; weight_decay replaces the config’s, None keeping it.

NO_DECAY_PATTERNS

attribute source

NO_DECAY_PATTERNS = ('*/bias', '*/scale', '*norm/weight', '*/dt_bias')

What lm-engine’s no_weight_decay group holds (configs/param-groups/ mup.yml at 45b6b57b): biases, every norm’s weight (an RMSNorm keeps scale here, Mamba-2’s gated norm weight) and Mamba-2’s dt_bias.

mup_param_groups

function source

def mup_param_groups(width_multiplier: float) -> tuple[ParamGroup, ...]

lm-engine’s muP parameter groups (configs/param-groups/mup.yml at 45b6b57b), in its order: norms, biases and dt_bias without weight decay at the base rate; the token embeddings at the base rate with decay; everything else, the router, A_log, D and the conv taps included, at the base rate divided by width_multiplier (lm-engine’s m_width, the model’s logits_scaling).

param_labels

function source

def param_labels(groups: Sequence[ParamGroup])

The optax.multi_transform labeller: each leaf’s first matching group.

power_schedule

function source

def power_schedule(
    peak: float,
    warmup_steps: int,
    a: float,
    b: float,
    c: float = 1.0,
    decay_start: int | None = None,
    decay_end: int | None = None,
    end_value: float = 0.0,
) -> optax.Schedule

lm-engine’s power scheduler (optimization/lr_scheduler/power.py at 45b6b57b) with an optional linear tail.

Past the warmup the rate is min(peak, a * (step * c) ** b): the power law of the batch size and step, arXiv 2408.13359, capped at peak (the optimizer’s own rate there). The warmup rises linearly from zero to that value at warmup_steps. decay_start set follows the law to that step and then decays linearly to end_value at decay_end, which is how Rigel ends its run; lm-engine’s scheduler has no tail.

linear_schedule

function source

def linear_schedule(
    peak: float,
    warmup_steps: int,
    decay_start: int | None,
    decay_end: int,
    end_value: float = 0.0,
) -> optax.Schedule

lm-engine’s linear scheduler (lr_scheduler/linear.py at 45b6b57b): from zero to peak over the warmup, constant to decay_start (None: the warmup’s end), then linear to end_value at decay_end.

ScheduleBase

class source

class ScheduleBase

One learning-rate schedule’s record: its own fields and its optax schedule. Registered under dew.registry.schedules, so a run’s record names its kind and holds no field another schedule reads.

ScheduleBase.schedule

def schedule(steps: int) -> optax.Schedule

The rate at each update of a steps-update run.

Cosine

dataclass source

class Cosine(
    peak: float,
    warmup_steps: int = 10000,
    end: float = 0.0,
    init: float = 0.0,
    decay_steps: int | None = None,
)

Linear warmup from init to peak, cosine to end at decay_steps (None: the run’s end); optax.warmup_cosine_decay_schedule.

Cosine.schedule

def schedule(steps: int) -> optax.Schedule

PowerTail

dataclass source

class PowerTail(start: int, steps: int | None = None, end: float = 0.0)

A power schedule’s linear tail: from the law’s rate at start to end at step steps (None: the run’s end).

Power

dataclass source

class Power(
    peak: float,
    warmup_steps: int,
    a: float,
    b: float = -0.51,
    c: float = 1.0,
    tail: PowerTail | None = None,
)

lm-engine’s power scheduler, power_schedule: warmup, then min(peak, a * (step * c) ** b), and with a tail a linear decay after the law (Rigel’s last 29%). lm-engine’s examples take a = 4 * batch size and c = tokens per step.

Power.schedule

def schedule(steps: int) -> optax.Schedule

Linear

dataclass source

class Linear(
    peak: float,
    warmup_steps: int = 0,
    decay_start: int | None = None,
    decay_steps: int | None = None,
    end: float = 0.0,
)

lm-engine’s linear scheduler, linear_schedule: warmup to peak, constant to decay_start (None: the warmup’s end), linear to end at decay_steps (None: the run’s end).

Linear.schedule

def schedule(steps: int) -> optax.Schedule

learning_rate_schedule

function source

def learning_rate_schedule(config: OptimConfig, steps: int)

The rate config names: its schedule over a steps-update run, or the constant learning_rate when it names none.

build_optimizer

function source

def build_optimizer(config: OptimConfig, steps: int) -> optax.GradientTransformation

Build the solver a config describes, with its schedule, parameter groups and clipping.

steps is the run’s length, which a schedule decays over unless the config names its own end. param_groups runs one solver per group under optax.multi_transform, each on the schedule times its multiplier and with its own weight decay; the global-norm clip still reads every gradient together, before the groups split them.