diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index d1a30ab3e..000e75398 100644 --- a/param_decomp/SPEC.md +++ b/param_decomp/SPEC.md @@ -289,7 +289,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g | S17 | `faithfulness_loss` is the global mean of squared delta entries over all sites' parameters (Σ‖Δ‖² / Σ numel), recomputed from live V/U each step. | | S18 | Each training step consumes a fresh token batch (a pure function of `(seed, step)` for O(1) resume); the model embeds it internally. **AMENDED 2026-06-24** (Oli-approved): removed the prefix harvest (residual-start) — there is no separate prefix forward; `clean_output` / `read_activations` / `masked_output` take the token batch directly. | | S19 | Components gradients are global-norm-clipped at `0.01`, before the optimizer step. CI fn is unclipped (production). The clip coefficient uses torch's eps convention: `clip_coef = min(1, max_norm / (total_norm + 1e-6))` (`torch.nn.utils.clip_grad_norm_`). This `+1e-6` is canonical and matched JAX-side (#643); plain `optax.clip_by_global_norm` divides by `max(total_norm, max_norm)` with NO eps, which at `clip=0.01` (clip fires almost every step) gives a ~1e-4 relative component-grad difference each step — a real per-step deviation the JAX clip must avoid by reproducing the `+1e-6`. | -| S20 | Both main optimizers are AdamW, `wd=0`, betas `(0.9, 0.999)`, eps `1e-8`; LR cosine to `0.1×` start, no warmup, stepped per training step. Cosine convention is canonical-torch: `progress = (step − warmup) / (decay_steps − 1)`, so LR reaches `0.1×` at `step = steps − 1` (`schedule.py`). This is matched JAX-side (#642); plain `optax.cosine_decay_schedule(lr, steps, alpha=0.1)` divides by `steps` and reaches `0.1×` one update LATER (at `count = steps`), a genuine formula divergence of ~O(1/steps) ≈ 2.5e-6/step at 400k that the JAX schedule must avoid by using the `decay_steps − 1` denominator. Both main LRs are evaluated by `scheduled_value_traced` honoring the full `ScheduleConfig`; the canonical cosine-to-0.1-no-warmup shape is pinned by the lab conversion gate (`assert_canonical_algorithm_config`). | +| S20 | Both main optimizers are AdamW, `wd=0`, betas `(0.9, 0.999)`, eps `1e-8`; LR cosine to `0.1×` start, no warmup, stepped per training step. Cosine convention is canonical-torch: `progress = (step − warmup) / (decay_steps − 1)`, so LR reaches `0.1×` at `step = steps − 1` (`schedule.py`). This is matched JAX-side (#642); plain `optax.cosine_decay_schedule(lr, steps, alpha=0.1)` divides by `steps` and reaches `0.1×` one update LATER (at `count = steps`), a genuine formula divergence of ~O(1/steps) ≈ 2.5e-6/step at 400k that the JAX schedule must avoid by using the `decay_steps − 1` denominator. Both main LRs are evaluated by `scheduled_value_traced` honoring the full `ScheduleConfig`; the canonical cosine-to-0.1-no-warmup shape is pinned by the lab conversion gate (`assert_canonical_algorithm_config`). **AMENDED 2026-07-02** (config-gated, default off = oracle parity): `type: muon` on a main optimizer config swaps that group's update rule for `optax.contrib.muon` (Newton-Schulz-orthogonalized momentum for the group's matrix leaves, Adam fallback for the rest), under the same cosine schedule and the same S19 clip chain. Every config without `type: muon` deserializes to `type: adamw` and keeps the canonical trajectory bit-identical. **AMENDED 2026-07-11** (muon × ci-fn matrix experiment): the muon matrix-leaf labeling is per-group (`run_state.build_optimizers`) — V/U and the MLP CI fns use optax's default 2D rule; the chunkwise CI fn passes `MuonDimensionNumbers` marking its 3D leaves as chunk-batched matrices (trailing two axes orthogonalized) and its 2D bias stacks as Adam-fallback, since the default 2D rule would label that tree exactly backwards. **AMENDED 2026-07-12** (fast muon): `impl: stacked` on a muon config swaps the per-leaf NS for `muon_stacked.py` — same-shape leaves batched into one NS with the stack axis sharded over `(replicate, fsdp)`, making each orthogonalization device-local (GSPMD otherwise lowers per-leaf NS on the ÷N masters into per-iteration full-Gram all-reduces with the largest matmul replicated on every device — the measured 3.3× muon-ci step hit). Same `MuonState` pytree (checkpoints round-trip across impls); trajectories match `impl: optax` up to float reassociation (the D4 tolerance class). `ns_steps` (default 5) and `ns_dtype` (default float32; bfloat16 = the Kimi recipe, stacked-only, masters/momentum stay fp32 per N1) are config knobs. Default `impl: optax` keeps the 07-02/07-11 experiment arms' exact semantics. | | S21 | Faithfulness warmup (400 × AdamW lr `1e-3` on `faithfulness_loss` alone) precedes step 0; its optimizer is discarded. | | S22 | Checkpoints round-trip ALL trajectory state of §3 — including every persistent term's sources + SRC_STEP moments + step/schedule counters — such that a resumed run continues the same trajectory (modulo RNG streams and kernel nondeterminism, cf. D4). **SIGTERM lifecycle (decision):** a SLURM-delivered SIGTERM sets a flag that is serviced at the main train-step boundary (synchronous save of the completed step, then exit for requeue), inside the faith-warmup loop (clean exit WITHOUT a save — no valid checkpoint exists pre-step-0, and resume skips warmup whenever any checkpoint is present, so a partial step-0 save would resume as if fully warmed; the requeue redoes warmup), and inside the in-loop eval pass (abandon the partial pass unlogged, fall through to the step-boundary save). The **first-jit-compile window remains unserviced** — a SIGTERM there is honored only once compilation finishes and control reaches the next servicing point. Periodic `save_every` is the BACKSTOP guarantee for every window; the SIGTERM→save path is the low-latency fast path, not the sole guarantee (lore `jsp_sigterm_save_never_fired`: 6/6 historical preemptions fell back to periodic ckpts; warmup/eval servicing closes the two largest unserviced windows). | | S23 | A persistent source bundle feeds exactly ONE loss term. (The fused-backward S14′ unscaling divides that term's coeff out of the source gradient; a bundle shared across terms would make the division wrong.) | diff --git a/param_decomp/configs.py b/param_decomp/configs.py index 455375223..1ff2fec6b 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -583,7 +583,8 @@ class ArithmeticCIGridConfig(BaseConfig): # --------------------------------------------------------------------------- -class OptimizerConfig(BaseConfig): +class AdamWOptimizerConfig(BaseConfig): + type: Literal["adamw"] = "adamw" lr_schedule: ScheduleConfig = Field(..., description="Learning rate schedule") weight_decay: NonNegativeFloat = Field(default=0.0, description="AdamW weight decay") betas: tuple[Probability, Probability] = Field( @@ -595,6 +596,69 @@ class OptimizerConfig(BaseConfig): ) +class MuonOptimizerConfig(BaseConfig): + """Muon (`optax.contrib.muon`): Newton-Schulz-orthogonalized momentum for the group's + matrix leaves; the rest fall back to Adam(0.9, 0.999) at the same LR. Experimental + (non-canonical). Which leaves are matrices is per-group (`run_state.build_optimizers`): + the V/U components tree is all-2D (fallback never fires); the chunkwise CI fn is + per-chunk stacks, so its 3D leaves are muon'd over the trailing two axes (chunk axis + batched) and its 2D bias stacks take the fallback; the MLP CI fns use the plain 2D rule.""" + + type: Literal["muon"] + lr_schedule: ScheduleConfig = Field(..., description="Learning rate schedule") + beta: Probability = Field( + default=0.95, description="Momentum decay for the orthogonalized update" + ) + consistent_rms: PositiveFloat | None = Field( + default=None, + description=( + "If set, scale updates by `sqrt(max(fan_in, fan_out)) * consistent_rms` so update" + " RMS is shape-independent (0.2 ~ AdamW's empirical RMS, making the AdamW LR" + " transferable). If None, optax's width scaling `sqrt(max(1, fan_out / fan_in))`." + ), + ) + weight_decay: NonNegativeFloat = Field(default=0.0, description="Weight decay") + grad_clip_norm: PositiveFloat | None = Field( + default=None, + description="If set, clip the grad norm of this group's parameters to this value", + ) + impl: Literal["optax", "stacked"] = Field( + default="optax", + description=( + "NS implementation. `optax` = per-leaf `optax.contrib.muon` (the 2026-07-02/11" + " experiment arms' exact semantics). `stacked` = same-shape leaves batched into" + " one NS with the stack axis sharded over (replicate, fsdp) — device-local" + " orthogonalization, no per-iteration collectives (`muon_stacked.py`); same" + " trajectory up to float reassociation (the SPEC D4 tolerance class)." + ), + ) + ns_steps: PositiveInt = Field( + default=5, description="Newton-Schulz iterations (optax default 5; fewer = cheaper/looser)" + ) + ns_dtype: Literal["float32", "bfloat16"] = Field( + default="float32", + description=( + "Dtype of the NS orthogonalization only (masters/momentum stay fp32 per N1);" + " bfloat16 halves NS compute+comm (the Kimi recipe). `stacked` impl only." + ), + ) + + +def _default_optimizer_type_adamw(data: object) -> object: + """AdamW is the canonical optimizer, so a config without `type` (every config predating + the muon gate, and the common case going forward) discriminates to it.""" + if isinstance(data, dict) and "type" not in data: + return {**data, "type": "adamw"} + return data + + +AnyOptimizerConfig = Annotated[ + AdamWOptimizerConfig | MuonOptimizerConfig, + Discriminator("type"), + BeforeValidator(_default_optimizer_type_adamw), +] + + AnyLossMetricConfig = Annotated[ ChunkwiseSubsetReconLossConfig | CIMaskedReconLayerwiseLossConfig @@ -933,10 +997,10 @@ def _strip_removed_jax_unsupported_fields(cls, data: object) -> object: ) # --- Training --- - components_optimizer: OptimizerConfig = Field( + components_optimizer: AnyOptimizerConfig = Field( ..., description="Optimizer config for the component (LinearComponent etc.) parameters" ) - ci_fn_optimizer: OptimizerConfig = Field( + ci_fn_optimizer: AnyOptimizerConfig = Field( ..., description="Optimizer config for the CI function parameters" ) steps: PositiveInt = Field(..., description="Total number of optimisation steps") diff --git a/param_decomp/muon_stacked.py b/param_decomp/muon_stacked.py new file mode 100644 index 000000000..f31a37912 --- /dev/null +++ b/param_decomp/muon_stacked.py @@ -0,0 +1,196 @@ +"""Stacked-NS Muon: optax.contrib.muon semantics with the Newton-Schulz orthogonalization +batched across same-shape leaves and sharded on the stack axis (SPEC S20, `impl: stacked`). + +Why: GSPMD lowers per-leaf NS on ÷N-sharded fp32 masters into per-iteration full-Gram +all-reduces with the largest matmul replicated on every device (~24.6 GB of serialized +collectives per step for the 4L ci-fn group — the measured 3.3x muon-ci step-time hit). +Stacking same-shape matrices and sharding the STACK axis makes each NS device-local: +one reshard in, one out, zero per-iteration collectives (the Kimi "Muon is Scalable" +parameter-partitioned recipe, expressed GSPMD-natively). + +Structure mirrors `optax.contrib.muon` exactly — same `MuonState(count, mu, ns_coeffs)`, +same muon/adam partition, same chain (NS -> consistent-rms shape scale -> weight decay -> +lr) — so checkpoints round-trip across `impl: optax|stacked` unchanged. Only the NS call +differs: leaves are canonicalized to `[g, rows<=cols]`, grouped by 2D shape, concatenated, +optionally cast to `ns_dtype`, orthogonalized once per group, and unstacked. Stacking +reorders float ops (concat + shared vmap), so trajectories match `impl: optax` only up to +reassociation — same tolerance class as device-count invariance (SPEC D4). +""" + +from collections import defaultdict +from collections.abc import Callable + +import jax +import jax.numpy as jnp +import optax +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P +from jaxtyping import Array + +# Not re-exported from `optax.contrib` in the pinned 0.2.8; the venv is pinned, so the +# private import is stable. `muon()` itself builds from these same pieces. +from optax.contrib._muon import orthogonalize_via_newton_schulz, scale_by_shape + +_NS_DIMS = optax.contrib.MuonDimensionNumbers(reduction_axis=-2, output_axis=-1) +_2D_DIMS = optax.contrib.MuonDimensionNumbers(reduction_axis=0, output_axis=1) + + +def _canonicalize(leaf: Array) -> tuple[Array, bool]: + """View a muon leaf as a `[g, a, b]` stack with `a <= b` (2D leaves get `g=1`). + NS orthogonalization and the `consistent_rms`/`sqrt(max(fan_in, fan_out))` scalings + are transpose-symmetric, so orientation only affects grouping, not semantics.""" + assert leaf.ndim in (2, 3), f"muon leaf must be 2D or a 3D stack, got {leaf.shape}" + stacked = leaf[None] if leaf.ndim == 2 else leaf + transposed = stacked.shape[-2] > stacked.shape[-1] + return (jnp.swapaxes(stacked, -2, -1) if transposed else stacked), transposed + + +def _pad_to_multiple(stack: Array, multiple: int) -> Array: + remainder = stack.shape[0] % multiple + if remainder == 0: + return stack + pad = multiple - remainder + return jnp.concatenate([stack, jnp.zeros((pad, *stack.shape[1:]), stack.dtype)], axis=0) + + +def _grouped_newton_schulz( + mu_hat: optax.Updates, + ns_coeffs: Array, + ns_steps: int, + ns_dtype: jnp.dtype, + stack_sharding: NamedSharding | None, +) -> optax.Updates: + """Orthogonalize every leaf of `mu_hat` (all muon-labeled by the partition) via ONE + batched NS per distinct canonical 2D shape, the batch axis sharded per + `stack_sharding` so each 2D orthogonalization runs device-local.""" + leaves, treedef = jax.tree.flatten(mu_hat) + canon = [_canonicalize(leaf) for leaf in leaves] + groups: dict[tuple[int, int], list[int]] = defaultdict(list) + for idx, (stacked, _) in enumerate(canon): + rows, cols = stacked.shape[-2:] + groups[(rows, cols)].append(idx) + + n_shards = 1 if stack_sharding is None else stack_sharding.num_devices + out: dict[int, Array] = {} + for shape, indices in groups.items(): + stacks = [canon[i][0] for i in indices] + sizes = [s.shape[0] for s in stacks] + out_dtype = stacks[0].dtype + # Cast BEFORE the sharding constraint so the ingress reshard moves ns_dtype + # bytes (half, for bf16 NS), not fp32. + grouped = jnp.concatenate(stacks, axis=0).astype(ns_dtype) + grouped = _pad_to_multiple(grouped, n_shards) + if stack_sharding is not None: + grouped = jax.lax.with_sharding_constraint(grouped, stack_sharding) + orthogonalized = orthogonalize_via_newton_schulz( + grouped, + jnp.asarray(ns_coeffs, ns_dtype), + ns_steps=ns_steps, + dimension_numbers=_NS_DIMS, + ).astype(out_dtype) + offset = 0 + for i, size in zip(indices, sizes, strict=True): + piece = orthogonalized[offset : offset + size] + offset += size + _, transposed = canon[i] + piece = jnp.swapaxes(piece, -2, -1) if transposed else piece + out[i] = piece[0] if leaves[i].ndim == 2 else piece + del shape + return jax.tree.unflatten(treedef, [out[i] for i in range(len(leaves))]) + + +def scale_by_stacked_muon( + *, + beta: float, + ns_steps: int, + ns_dtype: jnp.dtype, + stack_sharding: NamedSharding | None, +) -> optax.GradientTransformation: + """`optax.contrib.scale_by_muon` (nesterov, non-adaptive, frobenius) with the per-leaf + NS replaced by `_grouped_newton_schulz`. State is optax's `MuonState` verbatim.""" + reference = optax.contrib.scale_by_muon(beta=beta, ns_steps=ns_steps, nesterov=True) + + def update_fn( + updates: optax.Updates, state: optax.OptState, params: optax.Params | None = None + ) -> tuple[optax.Updates, optax.OptState]: + del params + assert isinstance(state, optax.contrib.MuonState) + mu = optax.tree.update_moment(updates, state.mu, beta, 1) + count_inc = optax.safe_increment(state.count) + mu_hat = jax.tree.map( + lambda m, g: beta * m + (1 - beta) * g, + optax.tree.bias_correction(mu, beta, optax.safe_increment(count_inc)), + optax.tree.bias_correction(updates, beta, count_inc), + ) + orthogonalized = _grouped_newton_schulz( + mu_hat, jnp.asarray(state.ns_coeffs), ns_steps, ns_dtype, stack_sharding + ) + return orthogonalized, optax.contrib.MuonState( + count=count_inc, mu=mu, ns_coeffs=state.ns_coeffs + ) + + return optax.GradientTransformation(reference.init, update_fn) + + +def stacked_muon( + learning_rate: optax.ScalarOrSchedule, + *, + beta: float, + weight_decay: float, + consistent_rms: float | None, + muon_weight_dimension_numbers: Callable[[optax.Params], optax.Params] | None, + ns_steps: int, + ns_dtype: jnp.dtype, + mesh: Mesh | None, +) -> optax.GradientTransformation: + """Drop-in for `optax.contrib.muon(...)` at our call site (`run_state`): same + muon/adam leaf partition, same post-NS chain, same state pytree — NS runs stacked, + sharded over `(replicate, fsdp)` when a mesh is given (None => single-device, e.g. + CPU tests and the toys).""" + stack_sharding = ( + NamedSharding(mesh, P(("replicate", "fsdp"), None, None)) if mesh is not None else None + ) + dim_nums = muon_weight_dimension_numbers + if dim_nums is None: + dim_nums = lambda params: jax.tree.map(lambda x: _NS_DIMS if x.ndim == 2 else None, params) + + def param_labels(params: optax.Params) -> optax.Params: + resolved = dim_nums(params) + return jax.tree.map( + lambda spec, x: jax.tree.map(lambda _: "muon" if spec is not None else "adam", x), + resolved, + params, + is_leaf=lambda x: x is None or isinstance(x, optax.contrib.MuonDimensionNumbers), + ) + + return optax.partition( + transforms={ + "muon": optax.chain( + scale_by_stacked_muon( + beta=beta, + ns_steps=ns_steps, + ns_dtype=ns_dtype, + stack_sharding=stack_sharding, + ), + scale_by_shape( + weight_dimension_numbers=lambda updates: jax.tree.map( + lambda x: _NS_DIMS if x.ndim == 3 else _2D_DIMS, updates + ), + consistent_rms=consistent_rms, + ), + optax.add_decayed_weights(weight_decay), + optax.scale_by_learning_rate(learning_rate), + ), + # Matches optax muon's fallback exactly: adamw at the muon lr, muon's + # nesterov=True threaded through (optax.adamw defaults nesterov=False). + "adam": optax.adamw( + learning_rate=learning_rate, + b1=0.9, + b2=0.999, + eps=1e-8, + weight_decay=0.0, + nesterov=True, + ), + }, + param_labels=param_labels, + ) diff --git a/param_decomp/run.py b/param_decomp/run.py index 17ebca0eb..52da863ae 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -490,7 +490,7 @@ def run_decomposition_training( save_every = cadence.save_every run.run_dir.mkdir(parents=True, exist_ok=True) - opt_vu, opt_ci, (sched_vu, sched_ci) = build_optimizers(pd) + opt_vu, opt_ci, (sched_vu, sched_ci) = build_optimizers(pd, mesh) key = random.PRNGKey(pd.seed) init_key, src_key, run_key = random.split(key, 3) diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index a20bfc436..ad4aa355a 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -21,9 +21,16 @@ from param_decomp.adversary import PersistentAdversary, init_sources_adam_state from param_decomp.built_run import DataConfig from param_decomp.ci_fn import CIFnArch -from param_decomp.configs import AdamPGDConfig, OptimizerConfig, PDConfig +from param_decomp.configs import ( + AdamPGDConfig, + AdamWOptimizerConfig, + ChunkwiseTransformerCiConfig, + MuonOptimizerConfig, + PDConfig, +) from param_decomp.lm import DecomposedModel from param_decomp.losses import scheduled_value_traced +from param_decomp.muon_stacked import stacked_muon from param_decomp.recon import ( PersistentSources, build_loss_terms, @@ -71,30 +78,80 @@ def update( return optax.GradientTransformation(init, update) -def _adamw_with_clip(opt: OptimizerConfig, schedule: Callable[[ArrayLike], Array]): - """AdamW (fp32 master, optax wd default overridden to the config's — torch's is 0) - over `schedule`, optionally preceded by torch-parity global-norm clip (SPEC S19/N1). - Adam eps is the torch/optax default 1e-8 (not exposed on `OptimizerConfig`).""" - adamw = optax.adamw( - schedule, b1=opt.betas[0], b2=opt.betas[1], eps=1e-8, weight_decay=opt.weight_decay - ) +def chunk_stacked_muon_dimension_numbers(params: optax.Params) -> optax.Params: + """Muon leaf labeling for the chunkwise CI fn's per-chunk stacks (`ci_fn.py`): every + 3D leaf is a `[n_chunks, d_in, d_out]` stack of `x @ W` matrices — orthogonalize the + trailing two axes, chunk axis batched — and everything else (`[n_chunks, d]` bias + stacks) takes the Adam fallback. optax's default rule (2D → muon) would do the + REVERSE on this tree: NS-orthogonalize the bias stacks and Adam the matrices.""" + dims = optax.contrib.MuonDimensionNumbers(reduction_axis=-2, output_axis=-1) + return jax.tree.map(lambda leaf: dims if leaf.ndim == 3 else None, params) + + +def _optimizer_with_clip( + opt: AdamWOptimizerConfig | MuonOptimizerConfig, + schedule: Callable[[ArrayLike], Array], + muon_dimension_numbers: Callable[[optax.Params], optax.Params] | None, + mesh: Mesh | None, +): + """The group optimizer (fp32 master) over `schedule`, optionally preceded by + torch-parity global-norm clip (SPEC S19/N1). AdamW is canonical (eps is the torch/optax + default 1e-8, not exposed on `AdamWOptimizerConfig`; optax's wd default overridden to the + config's — torch's is 0); Muon is a config-gated experimental variant (SPEC S19'). + `muon_dimension_numbers` labels the group's leaves for muon (None = optax's default + 2D-matrix rule, correct for the all-2D V/U tree and the MLP CI fns); ignored for adamw. + `mesh` shards the stacked-impl NS batch axis; None (toys, CPU tests) = unsharded.""" + match opt: + case AdamWOptimizerConfig(): + inner = optax.adamw( + schedule, b1=opt.betas[0], b2=opt.betas[1], eps=1e-8, weight_decay=opt.weight_decay + ) + case MuonOptimizerConfig(impl="optax"): + assert opt.ns_dtype == "float32", "ns_dtype is a stacked-impl knob (optax NS is fp32)" + inner = optax.contrib.muon( + schedule, + beta=opt.beta, + weight_decay=opt.weight_decay, + consistent_rms=opt.consistent_rms, + muon_weight_dimension_numbers=muon_dimension_numbers, + ns_steps=opt.ns_steps, + ) + case MuonOptimizerConfig(): + assert opt.impl == "stacked", opt.impl + inner = stacked_muon( + schedule, + beta=opt.beta, + weight_decay=opt.weight_decay, + consistent_rms=opt.consistent_rms, + muon_weight_dimension_numbers=muon_dimension_numbers, + ns_steps=opt.ns_steps, + ns_dtype=jnp.dtype(opt.ns_dtype), + mesh=mesh, + ) if opt.grad_clip_norm is None: - return adamw - return optax.chain(clip_by_global_norm_with_eps(opt.grad_clip_norm, eps=1e-6), adamw) + return inner + return optax.chain(clip_by_global_norm_with_eps(opt.grad_clip_norm, eps=1e-6), inner) -def build_optimizers(pd: PDConfig): +def build_optimizers(pd: PDConfig, mesh: Mesh | None): """Returns (opt_vu, opt_ci, schedules): the schedule fns are returned too so the log path reports the exact LR the optimizer applies (single source of truth). - The canonical-shape asserts (cosine-to-0.1, plain AdamW, required components clip, optional - CI-fn clip) live in + The canonical-shape asserts (cosine-to-0.1, canonical optimizer shape, required components + clip, optional CI-fn clip) live in the lab conversion (`experiments.config.assert_canonical_algorithm_config`); here we read the values straight off `PDConfig` so there is no second source of truth.""" sched_vu = optax_schedule(pd.components_optimizer.lr_schedule, pd.steps) sched_ci = optax_schedule(pd.ci_fn_optimizer.lr_schedule, pd.steps) - opt_vu = _adamw_with_clip(pd.components_optimizer, sched_vu) - opt_ci = _adamw_with_clip(pd.ci_fn_optimizer, sched_ci) + opt_vu = _optimizer_with_clip( + pd.components_optimizer, sched_vu, muon_dimension_numbers=None, mesh=mesh + ) + ci_muon_dim_nums = ( + chunk_stacked_muon_dimension_numbers + if isinstance(pd.ci_config, ChunkwiseTransformerCiConfig) + else None + ) + opt_ci = _optimizer_with_clip(pd.ci_fn_optimizer, sched_ci, ci_muon_dim_nums, mesh=mesh) return opt_vu, opt_ci, (sched_vu, sched_ci) diff --git a/param_decomp/tests/test_checkpoint.py b/param_decomp/tests/test_checkpoint.py index ea4cddbf2..556663954 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -2,6 +2,7 @@ trainer state (SPEC S22): a restored `TrainState` must continue the EXACT trajectory — including the persistent adversary's sources and Adam moments.""" +from collections.abc import Callable from pathlib import Path import equinox as eqx @@ -38,7 +39,9 @@ UniformKSubsetRoutingConfig, ) from param_decomp.lm import DecomposedModel +from param_decomp.muon_stacked import stacked_muon from param_decomp.recon import build_loss_terms +from param_decomp.run_state import chunk_stacked_muon_dimension_numbers from param_decomp.schedule import ScheduleConfig from param_decomp.sharding import hsdp_mesh from param_decomp.targets.llama8b import ( @@ -94,15 +97,37 @@ def _chunkwise_arch(lm: DecomposedModel, cfg: LlamaConfig) -> ChunkwiseTransform ) -def _build(seed: int): +def _build( + seed: int, muon_components: bool = False, muon_ci_fn: bool = False, stacked_impl: bool = False +): cfg = _tiny_cfg() C, seq = 8, 16 sites = llama_site_specs(cfg, mlp_family_site_cs(3, 4, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu(sites, jax.random.PRNGKey(seed)) ci_fn = build_ci_fn(_chunkwise_arch(lm, cfg), lm.sites, jax.random.PRNGKey(seed + 1)) - opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0)) - opt_ci = optax.adamw(1e-3, weight_decay=0.0) + + def muon_impl(dim_nums: "Callable[[optax.Params], optax.Params] | None"): + if stacked_impl: + return stacked_muon( + 1e-3, + beta=0.95, + weight_decay=0.0, + consistent_rms=0.2, + muon_weight_dimension_numbers=dim_nums, + ns_steps=5, + ns_dtype=jnp.dtype(jnp.float32), + mesh=None, + ) + return optax.contrib.muon(1e-3, consistent_rms=0.2, muon_weight_dimension_numbers=dim_nums) + + inner_vu = muon_impl(None) if muon_components else optax.adamw(1e-3, weight_decay=0.0) + opt_vu = optax.chain(optax.clip_by_global_norm(0.01), inner_vu) + opt_ci = ( + muon_impl(chunk_stacked_muon_dimension_numbers) + if muon_ci_fn + else optax.adamw(1e-3, weight_decay=0.0) + ) src = init_persistent_sources( lm.site_names, tuple(s.C for s in lm.sites), @@ -140,8 +165,12 @@ def _build(seed: int): return lm, state, step, resid -def test_roundtrip_and_exact_resume(tmp_path: Path): - lm, state, step, resid = _build(seed=1) +def _roundtrip_and_exact_resume( + tmp_path: Path, muon_components: bool, muon_ci_fn: bool = False, stacked_impl: bool = False +) -> None: + lm, state, step, resid = _build( + seed=1, muon_components=muon_components, muon_ci_fn=muon_ci_fn, stacked_impl=stacked_impl + ) for i in range(2): state, _ = step(lm, state, resid, jax.random.PRNGKey(i)) @@ -149,7 +178,9 @@ def test_roundtrip_and_exact_resume(tmp_path: Path): save_state(mgr, 2, state) # Restore onto a DIFFERENTLY-seeded reference: every leaf must come from disk. - _, fresh, _, _ = _build(seed=7) + _, fresh, _, _ = _build( + seed=7, muon_components=muon_components, muon_ci_fn=muon_ci_fn, stacked_impl=stacked_impl + ) restored = restore_latest(mgr, fresh) assert restored is not None loaded, ckpt_step = restored @@ -166,6 +197,31 @@ def test_roundtrip_and_exact_resume(tmp_path: Path): assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) +def test_roundtrip_and_exact_resume(tmp_path: Path): + _roundtrip_and_exact_resume(tmp_path, muon_components=False) + + +def test_muon_roundtrip_and_exact_resume(tmp_path: Path): + """SPEC S20 amendment: the muon components opt state (optax-partitioned muon/adam + masked trees) must ALSO restore onto a rebuilt reference and continue exactly — + this is what a scavenge preemption + requeue exercises.""" + _roundtrip_and_exact_resume(tmp_path, muon_components=True) + + +def test_muon_ci_fn_roundtrip_and_exact_resume(tmp_path: Path): + """SPEC S20 amendment (2026-07-11): same guarantee with muon on BOTH groups, the ci-fn + partitioned by `chunk_stacked_muon_dimension_numbers` (3D chunk stacks muon'd, 2D bias + stacks in the Adam-fallback mask).""" + _roundtrip_and_exact_resume(tmp_path, muon_components=True, muon_ci_fn=True) + + +def test_stacked_muon_roundtrip_and_exact_resume(tmp_path: Path): + """SPEC S20 `impl: stacked`: the stacked-NS muon state is optax's `MuonState` pytree + verbatim, so the same roundtrip + exact-resume guarantee holds — and a checkpoint + written under either impl restores under the other.""" + _roundtrip_and_exact_resume(tmp_path, muon_components=True, muon_ci_fn=True, stacked_impl=True) + + def test_persistent_adam_step_count_roundtrip_and_post_resume_bias_correction(tmp_path: Path): """Issue #678 (matrix §8 + S22/S13/S23): after N persistent ascents, the orbax checkpoint must carry the adversary's `step_count` leaf (present, fp32, == N) and diff --git a/param_decomp/tests/test_optim_torch_parity.py b/param_decomp/tests/test_optim_torch_parity.py index d5a5ccfb9..c0d6f4072 100644 --- a/param_decomp/tests/test_optim_torch_parity.py +++ b/param_decomp/tests/test_optim_torch_parity.py @@ -12,8 +12,15 @@ import pytest from jax.typing import ArrayLike from jaxtyping import Array +from pydantic import TypeAdapter -from param_decomp.run_state import clip_by_global_norm_with_eps, optax_schedule +from param_decomp.configs import AdamWOptimizerConfig, AnyOptimizerConfig, MuonOptimizerConfig +from param_decomp.run_state import ( + _optimizer_with_clip, + chunk_stacked_muon_dimension_numbers, + clip_by_global_norm_with_eps, + optax_schedule, +) from param_decomp.schedule import ScheduleConfig @@ -96,3 +103,183 @@ def test_grad_clip_noop_below_threshold(): out = _clip(clip, grads) assert float(out["a"][0]) == pytest.approx(3.0) assert float(out["a"][1]) == pytest.approx(4.0) + + +def test_muon_orthogonalizes_2d_leaves_and_adam_falls_back_elsewhere(): + """`type: muon` (SPEC S20 amendment): a 2D leaf's update is NS-orthogonalized (flat + singular values), a non-2D leaf falls back to Adam; default `type: adamw` keeps the + canonical optimizer so existing configs are untouched.""" + muon_cfg = MuonOptimizerConfig( + type="muon", + lr_schedule=ScheduleConfig( + fn_type="cosine", start_val=1e-3, final_val_frac=0.1, warmup_pct=0.0 + ), + grad_clip_norm=0.01, + ) + lr = 1e-3 + opt = _optimizer_with_clip(muon_cfg, lambda count: jnp.float32(lr), None, mesh=None) + key = jax.random.key(0) + params = {"V": jnp.zeros((16, 8)), "scale": jnp.zeros((8,))} + grads = { + "V": jax.random.normal(key, (16, 8)), + "scale": jax.random.normal(jax.random.fold_in(key, 1), (8,)), + } + updates, _ = opt.update(grads, opt.init(params), params) + _, treedef = jax.tree.flatten(grads) + updates = jax.tree.unflatten(treedef, jax.tree.leaves(updates)) + + grad_sv = jnp.linalg.svd(grads["V"], compute_uv=False) + update_sv = jnp.linalg.svd(updates["V"], compute_uv=False) + grad_flatness = float(grad_sv.max() / grad_sv.min()) + update_flatness = float(update_sv.max() / update_sv.min()) + assert update_flatness < 2.0 and update_flatness < grad_flatness / 2, ( + "muon update on a 2D leaf must be near-orthogonal (5-step NS is approximate, so the" + f" spectrum is flat-ish, not exactly flat): grad {grad_flatness:.2f} ->" + f" update {update_flatness:.2f}" + ) + scale_update_magnitude = float(jnp.abs(updates["scale"]).max()) + assert bool(jnp.all(jnp.isfinite(updates["scale"]))) + assert 0.3 * lr < scale_update_magnitude < 3 * lr, ( + f"non-2D leaf takes an Adam-fallback step of O(lr), got {scale_update_magnitude}" + ) + + +def test_muon_chunk_stacked_dimension_numbers_orthogonalize_3d_and_adam_2d_bias_stacks(): + """SPEC S20 amendment (2026-07-11): under `chunk_stacked_muon_dimension_numbers` a 3D + `[n_chunks, d_in, d_out]` matrix stack is NS-orthogonalized per chunk slice, while a 2D + `[n_chunks, d]` bias stack takes the Adam fallback — the reverse of optax's default 2D + rule, which on the chunkwise CI-fn tree would orthogonalize the bias stacks.""" + muon_cfg = MuonOptimizerConfig( + type="muon", + lr_schedule=ScheduleConfig( + fn_type="cosine", start_val=1e-3, final_val_frac=0.1, warmup_pct=0.0 + ), + grad_clip_norm=None, + ) + lr = 1e-3 + opt = _optimizer_with_clip( + muon_cfg, lambda count: jnp.float32(lr), chunk_stacked_muon_dimension_numbers, mesh=None + ) + key = jax.random.key(0) + n_chunks = 3 + params = {"w": jnp.zeros((n_chunks, 16, 8)), "b": jnp.zeros((n_chunks, 8))} + grads = { + "w": jax.random.normal(key, (n_chunks, 16, 8)), + "b": jax.random.normal(jax.random.fold_in(key, 1), (n_chunks, 8)), + } + updates, _ = opt.update(grads, opt.init(params), params) + _, treedef = jax.tree.flatten(grads) + updates = jax.tree.unflatten(treedef, jax.tree.leaves(updates)) + + for chunk in range(n_chunks): + grad_sv = jnp.linalg.svd(grads["w"][chunk], compute_uv=False) + update_sv = jnp.linalg.svd(updates["w"][chunk], compute_uv=False) + grad_flatness = float(grad_sv.max() / grad_sv.min()) + update_flatness = float(update_sv.max() / update_sv.min()) + assert update_flatness < 2.0 and update_flatness < grad_flatness / 2, ( + f"chunk {chunk}: 3D stack slice must be near-orthogonal:" + f" grad {grad_flatness:.2f} -> update {update_flatness:.2f}" + ) + bias_update_magnitude = float(jnp.abs(updates["b"]).max()) + assert bool(jnp.all(jnp.isfinite(updates["b"]))) + assert 0.3 * lr < bias_update_magnitude < 3 * lr, ( + f"2D bias stack takes an Adam-fallback step of O(lr), got {bias_update_magnitude}" + ) + + +def test_stacked_muon_update_matches_optax_muon(): + """SPEC S20 `impl: stacked`: the stack-by-shape batched NS produces the same updates as + per-leaf `optax.contrib.muon` (same momentum, same partition, same post-NS chain) up to + float reassociation — on a tree mixing 2D matrices (shared-shape group + a transposed + member), a 3D chunk stack, and Adam-fallback leaves.""" + schedule = ScheduleConfig(fn_type="cosine", start_val=1e-3, final_val_frac=0.1, warmup_pct=0.0) + lr = lambda count: jnp.float32(1e-3) + key = jax.random.key(7) + params = { + "w_a": jnp.zeros((16, 24)), + "w_b": jnp.zeros((24, 16)), + "stack": jnp.zeros((3, 16, 24)), + "bias_stack": jnp.zeros((3, 24)), + } + grads = { + name: jax.random.normal(jax.random.fold_in(key, i), p.shape) + for i, (name, p) in enumerate(params.items()) + } + dim_nums = chunk_stacked_muon_dimension_numbers + + from typing import Literal + + def cfg(impl: Literal["optax", "stacked"]) -> MuonOptimizerConfig: + return MuonOptimizerConfig( + type="muon", lr_schedule=schedule, grad_clip_norm=0.01, consistent_rms=0.2, impl=impl + ) + + def two_steps(impl: Literal["optax", "stacked"]): + opt = _optimizer_with_clip(cfg(impl), lr, dim_nums, mesh=None) + state = opt.init(params) + p = params + for _ in range(2): + updates, state = opt.update(grads, state, p) + p = jax.tree.map(lambda x, u: x + u, p, updates) + return p + + optax_p, stacked_p = two_steps("optax"), two_steps("stacked") + for k in params: + assert jnp.allclose(optax_p[k], stacked_p[k], rtol=1e-4, atol=1e-6), ( + f"{k}: stacked impl diverged from optax beyond reassociation tolerance" + ) + + +def test_stacked_muon_sharded_matches_unsharded(): + """The stack-axis sharding constraint is layout-only: at >1 devices the sharded + stacked NS reproduces the mesh=None updates (and preserves finiteness). Structural + no-op at 1 device; the real leg runs under + `XLA_FLAGS=--xla_force_host_platform_device_count=4`.""" + from jax.sharding import Mesh + + from param_decomp.muon_stacked import stacked_muon + from param_decomp.sharding import hsdp_mesh + + key = jax.random.key(3) + params = { + "w": jnp.zeros((4, 16, 24)), + "v": jnp.zeros((24, 16)), + "b": jnp.zeros((4, 24)), + } + grads = { + name: jax.random.normal(jax.random.fold_in(key, i), p.shape) + for i, (name, p) in enumerate(params.items()) + } + dim_nums = chunk_stacked_muon_dimension_numbers + + def one_update(mesh: "Mesh | None"): + opt = stacked_muon( + lambda count: jnp.float32(1e-3), + beta=0.95, + weight_decay=0.0, + consistent_rms=0.2, + muon_weight_dimension_numbers=dim_nums, + ns_steps=5, + ns_dtype=jnp.dtype(jnp.float32), + mesh=mesh, + ) + updates, _ = jax.jit(opt.update)(grads, opt.init(params), params) + _, treedef = jax.tree.flatten(grads) + return jax.tree.unflatten(treedef, jax.tree.leaves(updates)) + + unsharded, sharded = one_update(None), one_update(hsdp_mesh()) + for k in params: + assert bool(jnp.all(jnp.isfinite(sharded[k]))), k + assert jnp.allclose(unsharded[k], sharded[k], rtol=1e-5, atol=1e-7), ( + f"{k}: sharded stacked NS diverged from unsharded" + ) + + +def test_optimizer_config_type_discriminator(): + schedule = {"fn_type": "cosine", "start_val": 5e-5, "final_val_frac": 0.1} + adapter = TypeAdapter(AnyOptimizerConfig) + default = adapter.validate_python({"lr_schedule": schedule}) + assert isinstance(default, AdamWOptimizerConfig), "untyped configs stay canonical AdamW" + muon = adapter.validate_python({"type": "muon", "lr_schedule": schedule}) + assert isinstance(muon, MuonOptimizerConfig) + assert muon.beta == 0.95 and muon.consistent_rms is None diff --git a/param_decomp_lab/experiments/config.py b/param_decomp_lab/experiments/config.py index edc789793..e37501cc6 100644 --- a/param_decomp_lab/experiments/config.py +++ b/param_decomp_lab/experiments/config.py @@ -24,13 +24,14 @@ MLPCIArch, ) from param_decomp.configs import ( + AdamWOptimizerConfig, AnyEvalMetricConfig, Cadence, ChunkwiseTransformerCiConfig, CiConfig, GlobalMlpCiConfig, LayerwiseMlpCiConfig, - OptimizerConfig, + MuonOptimizerConfig, PDConfig, PersistentPGDReconLossConfig, ResumeProvenance, @@ -137,15 +138,21 @@ def _assert_cosine_to_tenth(schedule: ScheduleConfig, who: str) -> None: assert schedule.final_val_frac == 0.1, f"{who}: final_val_frac must be 0.1, got {schedule}" -def _assert_plain_adamw(optimizer: OptimizerConfig, who: str) -> None: - assert optimizer.betas == (0.9, 0.999), f"{who}: betas must be (0.9, 0.999)" +def _assert_canonical_optimizer( + optimizer: AdamWOptimizerConfig | MuonOptimizerConfig, who: str +) -> None: + """AdamW is canonical; Muon is a deliberate, config-gated experimental deviation (still + constrained to the trainer's no-weight-decay subspace).""" + if isinstance(optimizer, AdamWOptimizerConfig): + assert optimizer.betas == (0.9, 0.999), f"{who}: betas must be (0.9, 0.999)" assert optimizer.weight_decay == 0.0, f"{who}: weight_decay must be 0" def assert_canonical_algorithm_config(cfg: "ExperimentConfig[Any, Any]") -> None: """Assert the schema lives in the subspace the JAX trainer implements (the engine then reads `pd` / `cadence` DIRECTLY). The numerics-load-bearing constraints: - cosine-to-0.1 LR with no warmup, plain AdamW (betas (0.9, 0.999), no weight decay), + cosine-to-0.1 LR with no warmup, plain AdamW (betas (0.9, 0.999), no weight decay) or + the config-gated experimental Muon, a required components grad clip (CI-fn grad clip is optional), and a fully-specified checkpoint cadence. (Leaky-hard sigmoid, the always-built delta component, and no tied weights are now enforced by @@ -155,8 +162,8 @@ def assert_canonical_algorithm_config(cfg: "ExperimentConfig[Any, Any]") -> None ci_opt = cfg.pd.ci_fn_optimizer _assert_cosine_to_tenth(vu_opt.lr_schedule, "components_optimizer") _assert_cosine_to_tenth(ci_opt.lr_schedule, "ci_fn_optimizer") - _assert_plain_adamw(vu_opt, "components_optimizer") - _assert_plain_adamw(ci_opt, "ci_fn_optimizer") + _assert_canonical_optimizer(vu_opt, "components_optimizer") + _assert_canonical_optimizer(ci_opt, "ci_fn_optimizer") assert vu_opt.grad_clip_norm is not None, "components grad clip is part of the method" # The persistent-PGD source LR is constant-after-warmup only (`build_loss_terms` diff --git a/param_decomp_lab/experiments/lm/load_run.py b/param_decomp_lab/experiments/lm/load_run.py index c064d8591..589f9a9f5 100644 --- a/param_decomp_lab/experiments/lm/load_run.py +++ b/param_decomp_lab/experiments/lm/load_run.py @@ -155,7 +155,7 @@ def open_jax_run(run_dir: Path, step: int | None = None) -> LoadedJaxRun: mesh = hsdp_mesh() lm, vocab_size = build_target(cfg, mesh) - opt_vu, opt_ci, _ = build_optimizers(cfg.pd) + opt_vu, opt_ci, _ = build_optimizers(cfg.pd, mesh) init_key, src_key = jax.random.split(jax.random.PRNGKey(cfg.pd.seed)) reference = init_train_state( cfg.pd, lm, cfg.ci_fn, cfg.data, opt_vu, opt_ci, init_key, src_key, mesh