feat(optim): Muon — config-gated optimizer for both groups, with fast stacked-NS impl

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.) |

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")

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,
+    )

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)