kempnerforge.training.optimizer

Optimizer construction for KempnerForge.

Builds optimizers with per-parameter-group settings:

• AdamW: standard Adam with decoupled weight decay

• Muon: momentum with orthogonalized updates via Newton-Schulz iteration. Applies Muon to hidden nn.Linear weights, AdamW to the rest (embeddings, output head, norms, biases).

• Lion: sign-based momentum update (half the optimizer memory of AdamW)

• Schedule-Free AdamW: eliminates LR schedule by averaging iterates

All optimizers:

• Weight decay follows optimizer.weight_decay_scope: by default every parameter with ndim > 1 and no “bias” in its name is decayed, and “all” decays every parameter

• optimizer.lr_multipliers gives each listed module its own groups, at lr times the module’s multiplier

• Fused kernel when available (PyTorch 2.x, AdamW only)

Functions

build_optimizer(model, config[, module_patterns])	Construct an optimizer with per-parameter-group weight decay settings.
module_learning_rates(optimizer)	Return the rate each module build_optimizer gave its own groups steps at next.

Classes

Lion	Lion optimizer (Chen et al., 2023): Evolved Sign Momentum.
Muon	Muon optimizer: Momentum with Orthogonalized Updates.
ScheduleFreeAdamW	Schedule-Free AdamW (Defazio & Mishchenko, 2024).

class kempnerforge.training.optimizer.Lion

Bases: Optimizer

Lion optimizer (Chen et al., 2023): Evolved Sign Momentum.

Uses sign-based updates with momentum interpolation. Only maintains one momentum buffer (vs two for AdamW), halving optimizer memory.

Update rule:

update = sign(beta1 * m + (1 - beta1) * grad)
m = beta2 * m + (1 - beta2) * grad
p = p * (1 - lr * wd) - lr * update

Parameters:

• params – Parameters or parameter groups.

• lr – Learning rate (typically 3-10x smaller than AdamW).

• betas – (beta1, beta2) for update interpolation and momentum.

• weight_decay – Decoupled weight decay coefficient.

__init__(params, lr=0.0001, betas=(0.9, 0.99), weight_decay=0.0)

Parameters:

• lr (float)

• betas (tuple[float, float])

• weight_decay (float)

Return type:

None

step(closure=None)

class kempnerforge.training.optimizer.ScheduleFreeAdamW

Bases: Optimizer

Schedule-Free AdamW (Defazio & Mishchenko, 2024).

Eliminates the need for an LR schedule by maintaining an iterate z and a running average x. Parameters are set to the interpolated point y = (1 - beta1) * z + beta1 * x for gradient computation.

Use with scheduler.name = "none" — no LR schedule is needed.

Parameters:

• params – Parameters or parameter groups.

• lr – Learning rate (constant — no schedule needed).

• betas – (beta1, beta2) for interpolation and second moment.

• eps – Denominator term for numerical stability.

• weight_decay – Decoupled weight decay.

• warmup_steps – Linear warmup steps (internal to the optimizer).

__init__(params, lr=0.025, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, warmup_steps=0)

Parameters:

• lr (float)

• betas (tuple[float, float])

• eps (float)

• weight_decay (float)

• warmup_steps (int)

Return type:

None

state_dict()

Return type:

dict

load_state_dict(state_dict)

Parameters:

state_dict (dict)

Return type:

None

step(closure=None)

eval_params()

Set parameters to the evaluation point (running average).

Call before validation/inference for best results. Call train_params() afterward to resume training.

Return type:

None

train_params()

Restore parameters to the training point (interpolated y).

Call after eval_params() to resume training.

Return type:

None

class kempnerforge.training.optimizer.Muon

Bases: Optimizer

Muon optimizer: Momentum with Orthogonalized Updates.

For 2D+ weight matrices: maintains momentum, then orthogonalizes the update direction via Newton-Schulz iteration. This keeps update directions independent of parameter scale.

For the rest (embeddings, the output head, norms, biases): uses standard AdamW.

FSDP2 note: Newton-Schulz operates on each rank’s local shard independently — an approximation, not mathematically equivalent to orthogonalizing the full weight matrix. This is the standard approach for distributed Muon and works well in practice.

The AdamW fallback follows Muon’s schedule: each step sets every AdamW group’s rate to its base rate, its own lr or else adam_lr, times the ratio of Muon’s first group’s current rate to that group’s base rate.

Parameters:

• muon_params – Parameter groups for Muon (the hidden nn.Linear weights).

• adam_params – Parameter groups for AdamW fallback (every other parameter).

• lr – Learning rate for Muon.

• momentum – Momentum coefficient (default 0.95).

• weight_decay – Decoupled weight decay.

• adam_betas – Betas for the AdamW fallback.

• adam_eps – Epsilon for the AdamW fallback.

• ns_steps – Newton-Schulz iteration steps (default 5).

• adam_lr – Learning rate for AdamW fallback. None = same as lr.

• nesterov – Feed Newton-Schulz the Nesterov look-ahead instead of the momentum buffer.

• adjust_lr – Update scale: “original” sqrt(max(1, rows/cols)), “match_rms_adamw” 0.2 * sqrt(max(rows, cols)), or “spectral_unclamped” sqrt(rows/cols).

• orthogonalize – “local” runs Newton-Schulz on each rank’s shard, “exact” on the full matrix.

LEGACY_GROUP_DEFAULTS: ClassVar[dict[str, Any]] = {'adjust_lr': 'original', 'nesterov': False, 'orthogonalize': 'local'}

__init__(muon_params, adam_params, lr=0.02, momentum=0.95, weight_decay=0.0, adam_betas=(0.9, 0.95), adam_eps=1e-08, ns_steps=5, adam_lr=None, nesterov=False, adjust_lr='original', orthogonalize='local')

Parameters:

• muon_params (list[dict])

• adam_params (list[dict])

• lr (float)

• momentum (float)

• weight_decay (float)

• adam_betas (tuple[float, float])

• adam_eps (float)

• ns_steps (int)

• adam_lr (float | None)

• nesterov (bool)

• adjust_lr (str)

• orthogonalize (str)

property inner_optimizers: tuple[torch.optim.Optimizer, ...]

Optimizers holding parameters that Muon’s own param_groups do not list.

state_dict()

Include internal AdamW state in the plain state_dict(), and the base rates the fallback follows Muon’s schedule from when a param group has a rate of its own.

Return type:

dict

load_state_dict(state_dict)

Restore internal AdamW state from checkpoint.

A state without the fallback’s base rates keeps this optimizer’s own, and so does one whose rates do not number one per AdamW fallback group, with a warning.

Parameters:

state_dict (dict)

Return type:

None

zero_grad(set_to_none=True)

Zero both Muon’s own parameters and the inner AdamW’s.

Only the Muon half is registered with the base class, so the inherited zero_grad would leave every delegated parameter accumulating.

Parameters:

set_to_none (bool)

Return type:

None

step(closure=None)

kempnerforge.training.optimizer.module_learning_rates(optimizer)

Return the rate each module build_optimizer gave its own groups steps at next.

Keyed by the metric it is logged as, lr/<module>. Under Muon that is the rate of the module’s orthogonalized matrices, and lr_adamw/<module> is that of its AdamW fallback, each present when the module has such parameters. No key nests under another metric’s name, which a file-based tracking store cannot hold. Empty when optimizer.lr_multipliers was empty.

Parameters:

optimizer (torch.optim.Optimizer)

Return type:

dict[str, float]

kempnerforge.training.optimizer.build_optimizer(model, config, module_patterns=None)

Construct an optimizer with per-parameter-group weight decay settings.

With config.lr_multipliers set, each listed module’s parameters get their own groups at config.lr times the module’s multiplier, named by "module"; the scheduler then scales every group’s own rate by the same factor. A module nested in another listed one keeps its own parameters.

Parameters:

• model (torch.nn.Module) – Model whose parameters to optimize.

• config (OptimizerConfig) – Optimizer configuration.

• module_patterns (Mapping[str, Sequence[str]] | None) – Module name to the fnmatch patterns of its parameter names, such as VLMConfig.module_patterns. Required when config.lr_multipliers is set.

Returns:

Configured optimizer instance.

Raises:

ValueError – config.lr_multipliers names a module module_patterns lacks or one that matches no parameter, or two of its modules match the same parameters or parameters that partly overlap.

Return type:

torch.optim.Optimizer