From f423233fb5fa215dedbe010b23b1dca7a27e6c8b Mon Sep 17 00:00:00 2001 From: oli Date: Thu, 2 Jul 2026 10:23:42 +0000 Subject: [PATCH 01/13] fix(checkpoint): coerce restored state onto reference FORMAT not just sharding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prior device_put-onto-sharding was a confirmed no-op (StandardRestore already honors the sharding spec — verified via an instrumented resume: restored CI-fn was correctly ÷N). The real resume-OOM is a ÷1-scale ENTRY RELAYOUT on the first resumed step: orbax-restored arrays carry a default memory layout that differs from what the jitted step was compiled for, so the same executable that runs fresh at ~150GB materialized a ~103GB buffer on resume. The fresh-init reference is built by the same XLA layout assignment as the step, so its `.format` (layout + sharding) IS the step's expected input layout; coercing the restored tree onto it removes the entry relayout. Unblocks durable auto-resume (transient crashes can now recover). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR --- param_decomp/checkpoint.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/param_decomp/checkpoint.py b/param_decomp/checkpoint.py index 3223e02e2..927fff386 100644 --- a/param_decomp/checkpoint.py +++ b/param_decomp/checkpoint.py @@ -53,11 +53,13 @@ def restore_step(mgr: ocp.CheckpointManager, reference: TrainState, step: int) - (a freshly-initialised, correctly-placed `TrainState`).""" abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference) restored = mgr.restore(step, args=ocp.args.StandardRestore(abstract)) - # Force the restored tree onto the reference's exact shardings. StandardRestore honors the - # abstract target's shardings, but this makes the docstring's contract explicit and guards - # against a multi-host restore landing a leaf on a layout that differs from what the jitted - # step was compiled for (which triggers a costly entry-reshard on the first post-resume step). - restored = jax.device_put(restored, jax.tree.map(lambda s: s.sharding, abstract)) + # Coerce the restored tree onto the reference's exact FORMAT (layout + sharding), not just its + # sharding. StandardRestore already honors the sharding SPEC (verified), so a device_put onto + # sharding alone is a no-op — but orbax-restored arrays carry a default memory LAYOUT that + # differs from what the jitted step was compiled for. The reference is a fresh-init state built + # by the same XLA layout assignment as the step, so its `.format` IS the step's expected input + # layout; matching it avoids a ÷1-scale entry relayout on the first resumed step (the resume OOM). + restored = jax.device_put(restored, jax.tree.map(lambda r: r.format, reference)) return cast(TrainState, restored) From 74bfae945b864a09273a4d858e2af89a36acde77 Mon Sep 17 00:00:00 2001 From: Oliver Clive-Griffin Date: Thu, 2 Jul 2026 14:20:44 +0100 Subject: [PATCH 02/13] fix(hidden-acts-eval): derive waist leading from CI output, not token inputs (S31) (#919) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both hidden-acts eval steps computed leading = residual.shape[:-1] on what is actually the model INPUTS (an LM's [batch, seq] token ids), yielding (batch,) instead of (batch, seq). Every leading-shaped tensor (stochastic delta masks, zero deltas, the clean forward's all-false routes) lost the sequence axis and crashed site_out's route/delta broadcast the first time a run reached a slow eval with the metrics enabled — the deterministic CI step too, via route[..., None] in the clean forward, not just the stochastic delta path. leading now comes off the CI output [*leading, C] (_waist_leading), matching how lm.py::stochastic_site_masks already draws delta masks (ci.shape[:-1]) and staying target-generic. The misnamed residual arg is renamed to inputs: Any per the DecomposedModel protocol — the misnomer is what invited the bug; the sibling attn_patterns_eval.py names the same thing tokens. Regression test with batch != seq (fails pre-fix on both steps) pinning per-site sums/counts and the token-weighted accumulation. Claude-Session: https://claude.ai/code/session_016wiNqVkCMzQ8UQtkeEZ8eK Co-authored-by: Claude Fable 5 (cherry picked from commit ca215e74c53f47fdda6e318dd2b8c666ac681dd6) --- param_decomp/hidden_acts_eval.py | 46 ++++---- param_decomp/slow_eval.py | 9 +- param_decomp/tests/test_hidden_acts_eval.py | 115 ++++++++++++++++++++ 3 files changed, 147 insertions(+), 23 deletions(-) create mode 100644 param_decomp/tests/test_hidden_acts_eval.py diff --git a/param_decomp/hidden_acts_eval.py b/param_decomp/hidden_acts_eval.py index fd16df37e..99ce13cb2 100644 --- a/param_decomp/hidden_acts_eval.py +++ b/param_decomp/hidden_acts_eval.py @@ -29,7 +29,7 @@ import jax.numpy as jnp import numpy as np from jax import random -from jaxtyping import Array, Float, PRNGKeyArray +from jaxtyping import Array, PRNGKeyArray from param_decomp.components import DecompVU from param_decomp.jit_util import filter_jit @@ -57,13 +57,21 @@ def _per_site_sum_mse( HiddenActsStep = Callable[ - [DecomposedModel, Any, Any, Float[Array, "*leading d"], PRNGKeyArray], + [DecomposedModel, Any, Any, Any, PRNGKeyArray], tuple[dict[str, Array], dict[str, int]], ] -"""`(model, components, ci_fn, residual, key) -> ({site: sum_mse}, {site: n_elements})` +"""`(model, components, ci_fn, inputs, key) -> ({site: sum_mse}, {site: n_elements})` — one batch's per-site summed MSE (fp32) and element counts (the host folds these into -`SiteMSEReduction`s). `key` is unused by the deterministic CI step. `model` -(frozen-weight-bearing) is the jit ARG.""" +`SiteMSEReduction`s). `inputs` is the model's target-specific input (an LM's `[batch, seq]` +token ids), exactly as `read_activations` / `masked_site_outputs` take it. `key` is unused +by the deterministic CI step. `model` (frozen-weight-bearing) is the jit ARG.""" + + +def _waist_leading(ci_lower: dict[str, Array], site_names: tuple[str, ...]) -> tuple[int, ...]: + """The waist's `*leading` (batch + position axes), read off the CI output + `[*leading, C]` — the model `inputs` can't provide it (target-specific; an LM's token + ids carry no feature axis to strip).""" + return ci_lower[site_names[0]].shape[:-1] def make_ci_hidden_acts_step( @@ -76,24 +84,24 @@ def step( model: DecomposedModel, components: DecompVU, ci_fn: Any, - residual: Float[Array, "*leading d"], + inputs: Any, _key: PRNGKeyArray, ) -> tuple[dict[str, Array], dict[str, int]]: - taps = model.read_activations(residual, ci_fn.input_names) + taps = model.read_activations(inputs, ci_fn.input_names) components_bf16 = cast_floating(components, COMPUTE_DT) prepared = model.prepare_compute_weights(components_bf16) ci_fn_bf16 = cast_floating(ci_fn, COMPUTE_DT) ci_lower = ci_fn_bf16(taps, remat=False).lower - leading = residual.shape[:-1] + leading = _waist_leading(ci_lower, site_names) zeros_delta = {s: jnp.zeros(leading, COMPUTE_DT) for s in site_names} clean = model.masked_site_outputs( - prepared, residual, + prepared, inputs, {s: jnp.ones_like(ci_lower[s]) for s in site_names}, zeros_delta, all_false_routes(site_names, leading), site_names, False, ) # fmt: skip masked = model.masked_site_outputs( - prepared, residual, ci_lower, zeros_delta, None, site_names, False + prepared, inputs, ci_lower, zeros_delta, None, site_names, False ) sum_mse = _per_site_sum_mse(masked, clean, site_names) n_elements = {s: clean[s].size for s in site_names} @@ -117,18 +125,18 @@ def step( model: DecomposedModel, components: DecompVU, ci_fn: Any, - residual: Float[Array, "*leading d"], + inputs: Any, key: PRNGKeyArray, ) -> tuple[dict[str, Array], dict[str, int]]: - taps = model.read_activations(residual, ci_fn.input_names) + taps = model.read_activations(inputs, ci_fn.input_names) components_bf16 = cast_floating(components, COMPUTE_DT) prepared = model.prepare_compute_weights(components_bf16) ci_fn_bf16 = cast_floating(ci_fn, COMPUTE_DT) ci_lower = ci_fn_bf16(taps, remat=False).lower - leading = residual.shape[:-1] + leading = _waist_leading(ci_lower, site_names) clean = model.masked_site_outputs( - prepared, residual, + prepared, inputs, {s: jnp.ones_like(ci_lower[s]) for s in site_names}, {s: jnp.zeros(leading, COMPUTE_DT) for s in site_names}, all_false_routes(site_names, leading), site_names, False, @@ -148,7 +156,7 @@ def step( random.fold_in(delta_key, site_idx), leading, COMPUTE_DT ) masked = model.masked_site_outputs( - prepared, residual, masks, delta_masks, None, site_names, True + prepared, inputs, masks, delta_masks, None, site_names, True ) draw_sum = _per_site_sum_mse(masked, clean, site_names) sum_mse = {s: sum_mse[s] + draw_sum[s] for s in site_names} @@ -164,18 +172,18 @@ def accumulate_hidden_acts( model: DecomposedModel, components: DecompVU, ci_fn: Any, - residual_batches: list[Float[Array, "*leading d"]], + input_batches: list[Any], base_key: PRNGKeyArray, ) -> dict[str, SiteMSEReduction]: """Drive `step` over the eval batches, host-accumulating `(Σ sum_mse, Σ n)` per site (token-weighted, exact under micro-batching — same pattern as the density/mean accumulators). Per-batch RNG is `fold_in(base_key, batch_idx)`.""" - assert residual_batches, "hidden-acts eval needs at least one batch" + assert input_batches, "hidden-acts eval needs at least one batch" sums: dict[str, float] = {} counts: dict[str, int] = {} - for batch_idx, residual in enumerate(residual_batches): + for batch_idx, inputs in enumerate(input_batches): batch_sum, batch_n = step( - model, components, ci_fn, residual, random.fold_in(base_key, batch_idx) + model, components, ci_fn, inputs, random.fold_in(base_key, batch_idx) ) for site in batch_sum: sums[site] = sums.get(site, 0.0) + float(np.asarray(batch_sum[site])) diff --git a/param_decomp/slow_eval.py b/param_decomp/slow_eval.py index 219086792..f3301bed6 100644 --- a/param_decomp/slow_eval.py +++ b/param_decomp/slow_eval.py @@ -799,22 +799,23 @@ def render_slow_eval_figures( def compute_hidden_acts_metrics( model: DecomposedModel, state: Any, - residual_batches: list[Float[Array, "*leading d"]], + input_batches: list[Any], n_mask_samples: int, base_key: Array, compiler_options: dict[str, bool | int | str] | None = None, ) -> dict[str, float]: - """Both hidden-acts recon eval metrics over the eval batches, keyed by the torch + """Both hidden-acts recon eval metrics over the eval batches (`input_batches` holds + the model's target-specific inputs — an LM's token ids), keyed by the torch `[/]` log keys. `state.components`/`state.ci_fn` are the restored trajectory; `base_key` seeds the stochastic variant's per-batch draws.""" ci_key, stoch_key = random.split(base_key) ci_step = make_ci_hidden_acts_step(model, compiler_options) ci_reductions = accumulate_hidden_acts( - ci_step, model, state.components, state.ci_fn, residual_batches, ci_key + ci_step, model, state.components, state.ci_fn, input_batches, ci_key ) stoch_step = make_stochastic_hidden_acts_step(model, n_mask_samples, compiler_options) stoch_reductions = accumulate_hidden_acts( - stoch_step, model, state.components, state.ci_fn, residual_batches, stoch_key + stoch_step, model, state.components, state.ci_fn, input_batches, stoch_key ) return { **hidden_acts_log_entries("CIHiddenActsReconLoss", ci_reductions), diff --git a/param_decomp/tests/test_hidden_acts_eval.py b/param_decomp/tests/test_hidden_acts_eval.py new file mode 100644 index 000000000..c6d597ab8 --- /dev/null +++ b/param_decomp/tests/test_hidden_acts_eval.py @@ -0,0 +1,115 @@ +"""CPU tests for the in-loop hidden-acts recon eval metrics (SPEC S31). + +Pins the per-site MSE shape/count bookkeeping on both steps and the host-side +token-weighted accumulation. Batch != seq throughout: the steps must derive the waist +`*leading` from the CI output, not from the token inputs (whose `shape[:-1]` drops the +sequence axis — the regression that crashed the stochastic step's delta/route broadcast). +""" + +import jax +import numpy as np + +from param_decomp.ci_fn import Chunk, ChunkwiseTransformerCIArch, CIFn, build_ci_fn +from param_decomp.components import SiteC, init_decomp_vu +from param_decomp.hidden_acts_eval import ( + accumulate_hidden_acts, + hidden_acts_log_entries, + make_ci_hidden_acts_step, + make_stochastic_hidden_acts_step, +) +from param_decomp.lm import DecomposedModel +from param_decomp.targets.llama_simple_mlp import ( + canonical_site_cs, + parse_site_name, + site_specs, +) +from param_decomp.tests.test_llama_simple_mlp import ( + _tiny_cfg, + _tiny_decomposed_model, +) + +_BATCH, _SEQ = 2, 12 + + +def _build_ci_fn(lm: DecomposedModel, n_embd: int, key: jax.Array) -> CIFn: + site_names = lm.site_names + first_block = min(parse_site_name(n)[0] for n in site_names) + arch = ChunkwiseTransformerCIArch( + chunks=(Chunk(input_taps=(f"resid.{first_block}",), output_sites=site_names),), + input_dim=n_embd, + d_model=16, + n_blocks=1, + n_heads=2, + mlp_hidden=32, + ) + return build_ci_fn(arch, lm.sites, key) + + +def _setup(): + cfg = _tiny_cfg() + site_cs = canonical_site_cs( + ( + SiteC("h.2.attn.q_proj", 8), + SiteC("h.2.attn.v_proj", 12), + SiteC("h.2.mlp.c_fc", 8), + SiteC("h.3.mlp.down_proj", 16), + ) + ) + lm = _tiny_decomposed_model(cfg, site_specs(cfg, site_cs), jax.random.PRNGKey(0)) + components = init_decomp_vu(lm.sites, jax.random.PRNGKey(1)) + ci_fn = _build_ci_fn(lm, cfg.n_embd, jax.random.PRNGKey(2)) + tokens = jax.random.randint(jax.random.PRNGKey(3), (_BATCH, _SEQ), 0, cfg.vocab_size) + site_d_out = { + "h.2.attn.q_proj": cfg.n_head * cfg.head_dim, + "h.2.attn.v_proj": cfg.n_kv_head * cfg.head_dim, + "h.2.mlp.c_fc": cfg.n_intermediate, + "h.3.mlp.down_proj": cfg.n_embd, + } + return lm, components, ci_fn, tokens, site_d_out + + +def test_ci_step_per_site_sums_and_counts(): + lm, components, ci_fn, tokens, site_d_out = _setup() + step = make_ci_hidden_acts_step(lm) + + sum_mse, n_elements = step(lm, components, ci_fn, tokens, jax.random.PRNGKey(0)) + + assert set(sum_mse) == set(lm.site_names) == set(n_elements) + for site in lm.site_names: + assert int(n_elements[site]) == _BATCH * _SEQ * site_d_out[site] + assert np.isfinite(float(sum_mse[site])) + assert float(sum_mse[site]) >= 0.0 + + +def test_stochastic_step_per_site_sums_and_counts(): + lm, components, ci_fn, tokens, site_d_out = _setup() + n_mask_samples = 3 + step = make_stochastic_hidden_acts_step(lm, n_mask_samples) + + sum_mse, n_elements = step(lm, components, ci_fn, tokens, jax.random.PRNGKey(0)) + + assert set(sum_mse) == set(lm.site_names) == set(n_elements) + for site in lm.site_names: + assert int(n_elements[site]) == _BATCH * _SEQ * site_d_out[site] * n_mask_samples + assert np.isfinite(float(sum_mse[site])) + assert float(sum_mse[site]) >= 0.0 + + +def test_accumulate_and_log_entries_token_weighted(): + lm, components, ci_fn, tokens, _ = _setup() + step = make_ci_hidden_acts_step(lm) + + one = accumulate_hidden_acts(step, lm, components, ci_fn, [tokens], jax.random.PRNGKey(0)) + two = accumulate_hidden_acts( + step, lm, components, ci_fn, [tokens, tokens], jax.random.PRNGKey(0) + ) + + for site, r in two.items(): + assert r.n_elements == 2 * one[site].n_elements + np.testing.assert_allclose(r.sum_mse, 2 * one[site].sum_mse, rtol=1e-6) + + entries = hidden_acts_log_entries("CIHiddenActsReconLoss", two) + assert set(entries) == {"CIHiddenActsReconLoss"} | {f"CIHiddenActsReconLoss/{s}" for s in two} + total_sum = sum(r.sum_mse for r in two.values()) + total_n = sum(r.n_elements for r in two.values()) + np.testing.assert_allclose(entries["CIHiddenActsReconLoss"], total_sum / total_n) From 7ea29a104e4607d210a3de637689d247b49974f9 Mon Sep 17 00:00:00 2001 From: oli Date: Thu, 2 Jul 2026 17:39:52 +0000 Subject: [PATCH 03/13] feat(optim): config-gated Muon for the main optimizers (amends SPEC S20) - AnyOptimizerConfig: discriminated adamw|muon union; untyped configs default to adamw (canonical, bit-identical trajectories). - MuonOptimizerConfig: optax.contrib.muon (NS-orthogonalized momentum on 2D leaves, Adam fallback elsewhere), consistent_rms knob, same cosine schedule + S19 clip chain. - Lab canonical assert admits muon (still no-weight-decay subspace). - Tests: muon orthogonalization + adam fallback + discriminator (test_optim_torch_parity), muon checkpoint roundtrip + exact resume (test_checkpoint, S22). - configs/muon/: pile 4L PPGD 40k A/B arms (control + muon 1x/4x lr, consistent_rms 0.2). Co-Authored-By: Claude Fable 5 --- param_decomp/SPEC.md | 2 +- param_decomp/configs.py | 47 +++++- .../configs/muon/pile4l_ppgd_40k_control.yaml | 157 ++++++++++++++++++ .../configs/muon/pile4l_ppgd_40k_muon_1x.yaml | 157 ++++++++++++++++++ .../configs/muon/pile4l_ppgd_40k_muon_4x.yaml | 157 ++++++++++++++++++ param_decomp/run_state.py | 40 +++-- param_decomp/tests/test_checkpoint.py | 26 ++- param_decomp/tests/test_optim_torch_parity.py | 58 ++++++- param_decomp_lab/experiments/config.py | 19 ++- 9 files changed, 633 insertions(+), 30 deletions(-) create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_control.yaml create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index c009e4222..670f609ac 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. | +| 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. **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 2D 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. | | 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 781ea7b0f..f2c92e339 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -564,7 +564,8 @@ class UVPlotsConfig(_PermutationPlotsBaseConfig): # --------------------------------------------------------------------------- -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( @@ -576,6 +577,46 @@ class OptimizerConfig(BaseConfig): ) +class MuonOptimizerConfig(BaseConfig): + """Muon (`optax.contrib.muon`): Newton-Schulz-orthogonalized momentum for the group's 2D + leaves; non-2D leaves fall back to Adam(0.9, 0.999) at the same LR. Experimental + (non-canonical) — the V/U components tree is all-2D so the fallback never fires there.""" + + 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", + ) + + +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 @@ -916,10 +957,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/configs/muon/pile4l_ppgd_40k_control.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_control.yaml new file mode 100644 index 000000000..6e5737161 --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_control.yaml @@ -0,0 +1,157 @@ +# try-muon A/B (bridge/task-try-muon): byte-derived from baseline p-5d480c2d's pinned +# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this +# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix +# cherry-picked here). Control arm: canonical AdamW, trajectory-identical to stock. +run_name: jax-pile4l-ppgd-muon-control +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: try-muon + tags: + - try-muon diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml new file mode 100644 index 000000000..6a2a7d382 --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml @@ -0,0 +1,157 @@ +# try-muon A/B (bridge/task-try-muon): byte-derived from baseline p-5d480c2d's pinned +# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this +# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix +# cherry-picked here). Muon arm (1x = lr 5.0e-05): components_optimizer type=muon, consistent_rms 0.2 (update +# RMS matched to AdamW's per the Kimi recipe, so lr is on the AdamW scale), same cosine +# schedule + S19 clip; ci_fn_optimizer stays canonical AdamW for clean attribution. +run_name: jax-pile4l-ppgd-muon-1x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: try-muon + tags: + - try-muon diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml new file mode 100644 index 000000000..b667fc7be --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml @@ -0,0 +1,157 @@ +# try-muon A/B (bridge/task-try-muon): byte-derived from baseline p-5d480c2d's pinned +# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this +# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix +# cherry-picked here). Muon arm (4x = lr 2.0e-04): components_optimizer type=muon, consistent_rms 0.2 (update +# RMS matched to AdamW's per the Kimi recipe, so lr is on the AdamW scale), same cosine +# schedule + S19 clip; ci_fn_optimizer stays canonical AdamW for clean attribution. +run_name: jax-pile4l-ppgd-muon-4x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 2.0e-04 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: try-muon + tags: + - try-muon diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index 38e99af41..7bd989b18 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -21,7 +21,7 @@ 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, MuonOptimizerConfig, PDConfig from param_decomp.lm import DecomposedModel from param_decomp.recon import ( PersistentSources, @@ -80,32 +80,44 @@ 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 _optimizer_with_clip( + opt: AdamWOptimizerConfig | MuonOptimizerConfig, schedule: Callable[[ArrayLike], Array] +): + """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').""" + 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(): + inner = optax.contrib.muon( + schedule, + beta=opt.beta, + weight_decay=opt.weight_decay, + consistent_rms=opt.consistent_rms, + ) 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): """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 = torch_cosine_schedule( pd.components_optimizer.lr_schedule.start_val, pd.steps, alpha=0.1 ) sched_ci = torch_cosine_schedule(pd.ci_fn_optimizer.lr_schedule.start_val, pd.steps, alpha=0.1) - 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) + opt_ci = _optimizer_with_clip(pd.ci_fn_optimizer, sched_ci) 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 e50716656..eb47b91d8 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -94,14 +94,19 @@ def _chunkwise_arch(lm: DecomposedModel, cfg: LlamaConfig) -> ChunkwiseTransform ) -def _build(seed: int): +def _build(seed: int, muon_components: 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)) + inner_vu = ( + optax.contrib.muon(1e-3, consistent_rms=0.2) + 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 = optax.adamw(1e-3, weight_decay=0.0) src = init_persistent_sources( lm.site_names, @@ -141,8 +146,8 @@ 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) -> None: + lm, state, step, resid = _build(seed=1, muon_components=muon_components) for i in range(2): state, _ = step(lm, state, resid, jax.random.PRNGKey(i)) @@ -150,7 +155,7 @@ 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) restored = restore_latest(mgr, fresh) assert restored is not None loaded, ckpt_step = restored @@ -167,6 +172,17 @@ 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_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 efa6f8ff2..ffbb0645f 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, torch_cosine_schedule +from param_decomp.configs import AdamWOptimizerConfig, AnyOptimizerConfig, MuonOptimizerConfig +from param_decomp.run_state import ( + _optimizer_with_clip, + clip_by_global_norm_with_eps, + torch_cosine_schedule, +) +from param_decomp.schedule import ScheduleConfig def _scalar(value: ArrayLike) -> float: @@ -92,3 +99,52 @@ 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)) + 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_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 901facc3d..8dc2a2e1c 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, @@ -136,15 +137,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 @@ -154,8 +161,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 — `source_lr` ignores From 7741ccc76a05458a5148c1046db06159c07e661f Mon Sep 17 00:00:00 2001 From: oli Date: Fri, 3 Jul 2026 10:31:06 +0000 Subject: [PATCH 04/13] config(muon): adamw-4x-lr disambiguation control arm Co-Authored-By: Claude Fable 5 --- .../muon/pile4l_ppgd_40k_adamw_4x.yaml | 157 ++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml new file mode 100644 index 000000000..283769bd2 --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml @@ -0,0 +1,157 @@ +# try-muon A/B: adamw at 4x lr (2e-4) — the LR-vs-optimizer disambiguation control for muon_4x. +# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this +# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix +# cherry-picked here). Control arm: canonical AdamW, trajectory-identical to stock. +run_name: jax-pile4l-ppgd-adamw-4x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 2.0e-04 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: try-muon + tags: + - try-muon From af318391a644cebc3bc010a75d84bb9592e61c63 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sat, 11 Jul 2026 01:10:17 +0000 Subject: [PATCH 05/13] feat(optim): muon for the chunkwise ci-fn optimizer (MuonDimensionNumbers over chunk stacks) The chunkwise CI fn's trainable leaves are per-chunk stacks: 3D [n_chunks, d_in, d_out] matrix stacks and 2D [n_chunks, d] bias stacks. optax's default muon rule (2D -> muon, rest -> Adam) labels that tree exactly backwards, so build_optimizers now passes explicit MuonDimensionNumbers for the ci-fn group when the arch is chunkwise: 3D leaves are NS-orthogonalized over the trailing two axes (chunk axis batched), everything else takes the Adam fallback. V/U and the MLP CI fns keep the default 2D rule (all-2D / plain-2D trees, unchanged). SPEC S20 amendment extended (2026-07-11). Two matrix arm configs added, byte-derived from the try-muon control / muon-1x stored launch configs with only ci_fn_optimizer flipped to muon (lr 5e-5, consistent_rms 0.2): the two missing cells of the {adamw, muon} x {UV, ci-fn} matrix. Tests: muon 3D-stack orthogonalization + 2D-bias-stack Adam fallback (test_optim_torch_parity), muon-both-groups checkpoint roundtrip + exact resume (test_checkpoint, the scavenge-preemption path). Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- param_decomp/SPEC.md | 2 +- param_decomp/configs.py | 9 +- .../muon/pile4l_ppgd_40k_both_muon.yaml | 154 +++++++++++++++++ .../configs/muon/pile4l_ppgd_40k_ci_muon.yaml | 157 ++++++++++++++++++ param_decomp/run_state.py | 36 +++- param_decomp/tests/test_checkpoint.py | 28 +++- param_decomp/tests/test_optim_torch_parity.py | 46 ++++- 7 files changed, 417 insertions(+), 15 deletions(-) create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index 670f609ac..59b1aaa5f 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. **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 2D 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. | +| 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. **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. | | 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 f2c92e339..4b28d4c7c 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -578,9 +578,12 @@ class AdamWOptimizerConfig(BaseConfig): class MuonOptimizerConfig(BaseConfig): - """Muon (`optax.contrib.muon`): Newton-Schulz-orthogonalized momentum for the group's 2D - leaves; non-2D leaves fall back to Adam(0.9, 0.999) at the same LR. Experimental - (non-canonical) — the V/U components tree is all-2D so the fallback never fires there.""" + """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") diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml new file mode 100644 index 000000000..bd6e9de32 --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml @@ -0,0 +1,154 @@ +# muon x {UV, ci-fn} matrix (bridge/task-muon-cifn-matrix): byte-derived from the try-muon +# muon-1x arm p-cedbbf7c's config (= pile_ppgd_bsc.yaml, steps 400k->40k, dp 16->8, seed 0, +# components_optimizer muon @ lr 5e-5 consistent_rms 0.2), with ONLY ci_fn_optimizer ALSO +# flipped to muon (same settings, still unclipped per S19). Matrix cell: UV=muon, ci-fn=muon. +run_name: jax-pile4l-ppgd-both-muon +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml new file mode 100644 index 000000000..b9d0633a9 --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml @@ -0,0 +1,157 @@ +# muon x {UV, ci-fn} matrix (bridge/task-muon-cifn-matrix): byte-derived from the try-muon +# control arm p-d0f66d3b's config (= pile_ppgd_bsc.yaml, steps 400k->40k, dp 16->8, seed 0), +# with ONLY ci_fn_optimizer flipped to muon (consistent_rms 0.2, same lr/schedule, still +# unclipped per S19). Matrix cell: UV=adamw, ci-fn=muon. Chunk-stacked 3D ci-fn leaves are +# muon'd via chunk_stacked_muon_dimension_numbers; 2D bias stacks take the Adam fallback. +run_name: jax-pile4l-ppgd-ci-muon +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index 7bd989b18..7e93ad24a 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -21,7 +21,13 @@ 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, AdamWOptimizerConfig, MuonOptimizerConfig, PDConfig +from param_decomp.configs import ( + AdamPGDConfig, + AdamWOptimizerConfig, + ChunkwiseTransformerCiConfig, + MuonOptimizerConfig, + PDConfig, +) from param_decomp.lm import DecomposedModel from param_decomp.recon import ( PersistentSources, @@ -80,13 +86,27 @@ def update( return optax.GradientTransformation(init, update) +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] + opt: AdamWOptimizerConfig | MuonOptimizerConfig, + schedule: Callable[[ArrayLike], Array], + muon_dimension_numbers: Callable[[optax.Params], optax.Params] | 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').""" + 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.""" match opt: case AdamWOptimizerConfig(): inner = optax.adamw( @@ -98,6 +118,7 @@ def _optimizer_with_clip( beta=opt.beta, weight_decay=opt.weight_decay, consistent_rms=opt.consistent_rms, + muon_weight_dimension_numbers=muon_dimension_numbers, ) if opt.grad_clip_norm is None: return inner @@ -116,8 +137,13 @@ def build_optimizers(pd: PDConfig): pd.components_optimizer.lr_schedule.start_val, pd.steps, alpha=0.1 ) sched_ci = torch_cosine_schedule(pd.ci_fn_optimizer.lr_schedule.start_val, pd.steps, alpha=0.1) - opt_vu = _optimizer_with_clip(pd.components_optimizer, sched_vu) - opt_ci = _optimizer_with_clip(pd.ci_fn_optimizer, sched_ci) + opt_vu = _optimizer_with_clip(pd.components_optimizer, sched_vu, muon_dimension_numbers=None) + 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) 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 eb47b91d8..6aea1963b 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -39,6 +39,7 @@ ) from param_decomp.lm import DecomposedModel 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,7 +95,7 @@ def _chunkwise_arch(lm: DecomposedModel, cfg: LlamaConfig) -> ChunkwiseTransform ) -def _build(seed: int, muon_components: bool = False): +def _build(seed: int, muon_components: bool = False, muon_ci_fn: bool = False): cfg = _tiny_cfg() C, seq = 8, 16 sites = llama_site_specs(cfg, mlp_family_site_cs(3, 4, C)) @@ -107,7 +108,15 @@ def _build(seed: int, muon_components: bool = False): else optax.adamw(1e-3, weight_decay=0.0) ) opt_vu = optax.chain(optax.clip_by_global_norm(0.01), inner_vu) - opt_ci = optax.adamw(1e-3, weight_decay=0.0) + opt_ci = ( + optax.contrib.muon( + 1e-3, + consistent_rms=0.2, + muon_weight_dimension_numbers=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), @@ -146,8 +155,10 @@ def _build(seed: int, muon_components: bool = False): return lm, state, step, resid -def _roundtrip_and_exact_resume(tmp_path: Path, muon_components: bool) -> None: - lm, state, step, resid = _build(seed=1, muon_components=muon_components) +def _roundtrip_and_exact_resume( + tmp_path: Path, muon_components: bool, muon_ci_fn: bool = False +) -> None: + lm, state, step, resid = _build(seed=1, muon_components=muon_components, muon_ci_fn=muon_ci_fn) for i in range(2): state, _ = step(lm, state, resid, jax.random.PRNGKey(i)) @@ -155,7 +166,7 @@ def _roundtrip_and_exact_resume(tmp_path: Path, muon_components: bool) -> None: save_state(mgr, 2, state) # Restore onto a DIFFERENTLY-seeded reference: every leaf must come from disk. - _, fresh, _, _ = _build(seed=7, muon_components=muon_components) + _, fresh, _, _ = _build(seed=7, muon_components=muon_components, muon_ci_fn=muon_ci_fn) restored = restore_latest(mgr, fresh) assert restored is not None loaded, ckpt_step = restored @@ -183,6 +194,13 @@ def test_muon_roundtrip_and_exact_resume(tmp_path: Path): _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_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 ffbb0645f..b8e90c158 100644 --- a/param_decomp/tests/test_optim_torch_parity.py +++ b/param_decomp/tests/test_optim_torch_parity.py @@ -17,6 +17,7 @@ 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, torch_cosine_schedule, ) @@ -113,7 +114,7 @@ def test_muon_orthogonalizes_2d_leaves_and_adam_falls_back_elsewhere(): grad_clip_norm=0.01, ) lr = 1e-3 - opt = _optimizer_with_clip(muon_cfg, lambda count: jnp.float32(lr)) + opt = _optimizer_with_clip(muon_cfg, lambda count: jnp.float32(lr), None) key = jax.random.key(0) params = {"V": jnp.zeros((16, 8)), "scale": jnp.zeros((8,))} grads = { @@ -140,6 +141,49 @@ def test_muon_orthogonalizes_2d_leaves_and_adam_falls_back_elsewhere(): ) +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 + ) + 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_optimizer_config_type_discriminator(): schedule = {"fn_type": "cosine", "start_val": 5e-5, "final_val_frac": 0.1} adapter = TypeAdapter(AnyOptimizerConfig) From dde89eb00061e308756eb1d56009e4d166bf21f5 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sat, 11 Jul 2026 01:28:41 +0000 Subject: [PATCH 06/13] config(muon-matrix): disable XLA per-fusion autotune persistent cache for the matrix arms The shared cache's xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable on 2026-07-08, so any non-owner run dies PERMISSION_DENIED at the first train step (job 1052115). JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES=none skips only that side cache; the main executable cache still applies. The 2026-07-03 sibling arms predate the dir, so this also matches their autotune behavior. Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml | 7 +++++++ param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml index bd6e9de32..44127f39e 100644 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml @@ -139,6 +139,13 @@ pd: runtime: remat_recon_forwards: false dp: 8 + launch_env: + env: + # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + # so disabling it also matches their autotune behavior. + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none target: output_extract: 0 weights_dtype: bfloat16 diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml index b9d0633a9..bcaf0f19b 100644 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml @@ -142,6 +142,13 @@ pd: runtime: remat_recon_forwards: false dp: 8 + launch_env: + env: + # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + # so disabling it also matches their autotune behavior. + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none target: output_extract: 0 weights_dtype: bfloat16 From 3abf69411d93925f3a330f9d39f76886d572cdd9 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sat, 11 Jul 2026 08:23:48 +0000 Subject: [PATCH 07/13] =?UTF-8?q?config(muon-matrix):=20imp-min=20coeff=20?= =?UTF-8?q?sweep=20=E2=80=94=20each=20matrix=20cell=20at=202x=20and=200.5x?= =?UTF-8?q?=20(4e-4,=201e-4)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Oli's follow-up on the muon pilot thread: the 4 matrix cells land at different sparsity-recon tradeoff points, so compare optimizer setups as 3-point Pareto curves (L0 vs adv/stoch recon at step 30k) instead of single runs. Everything except run_name and the ImportanceMinimalityLoss coeff is byte-identical to the parent cell's config; full 40k schedule so the 30k point is comparable to the existing coeff-2e-4 runs. Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- .../muon/impmin_sweep/both_muon_imp05x.yaml | 159 +++++++++++++++++ .../muon/impmin_sweep/both_muon_imp2x.yaml | 159 +++++++++++++++++ .../muon/impmin_sweep/ci_muon_imp05x.yaml | 161 ++++++++++++++++++ .../muon/impmin_sweep/ci_muon_imp2x.yaml | 161 ++++++++++++++++++ .../muon/impmin_sweep/control_imp05x.yaml | 160 +++++++++++++++++ .../muon/impmin_sweep/control_imp2x.yaml | 160 +++++++++++++++++ .../muon/impmin_sweep/muon_1x_imp05x.yaml | 158 +++++++++++++++++ .../muon/impmin_sweep/muon_1x_imp2x.yaml | 158 +++++++++++++++++ 8 files changed, 1276 insertions(+) create mode 100644 param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml create mode 100644 param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml create mode 100644 param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml create mode 100644 param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml create mode 100644 param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml create mode 100644 param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml create mode 100644 param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml create mode 100644 param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml diff --git a/param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml new file mode 100644 index 000000000..e8a18a199 --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml @@ -0,0 +1,159 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): both_muon cell at imp-min +# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_both_muon.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-both-muon-imp05x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0001 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml new file mode 100644 index 000000000..a4ff9a1c8 --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml @@ -0,0 +1,159 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): both_muon cell at imp-min +# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_both_muon.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-both-muon-imp2x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0004 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml new file mode 100644 index 000000000..b73d3089b --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml @@ -0,0 +1,161 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): ci_muon cell at imp-min +# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_ci_muon.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-ci-muon-imp05x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0001 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml new file mode 100644 index 000000000..0dc50b6ba --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml @@ -0,0 +1,161 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): ci_muon cell at imp-min +# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_ci_muon.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-ci-muon-imp2x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0004 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml new file mode 100644 index 000000000..b16f4ae9c --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml @@ -0,0 +1,160 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): control cell at imp-min +# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_control.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-muon-control-imp05x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0001 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml new file mode 100644 index 000000000..651e1c7dc --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml @@ -0,0 +1,160 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): control cell at imp-min +# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_control.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-muon-control-imp2x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0004 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml new file mode 100644 index 000000000..4447394cb --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml @@ -0,0 +1,158 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): muon_1x cell at imp-min +# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_muon_1x.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-muon-1x-imp05x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0001 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml new file mode 100644 index 000000000..641f20963 --- /dev/null +++ b/param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml @@ -0,0 +1,158 @@ +# impmin sweep (bridge/task muon-cifn-matrix follow-up): muon_1x cell at imp-min +# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_muon_1x.yaml, only run_name / coeff / +# (env fix + tags for the try-muon-era bases) changed. +run_name: jax-pile4l-ppgd-muon-1x-imp2x +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0004 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: muon-cifn-matrix + tags: + - muon-cifn-matrix + - impmin-sweep From faf98fc7805aaf97880c5204ce8503a5af1dfa31 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sun, 12 Jul 2026 16:42:00 +0000 Subject: [PATCH 08/13] config(smooth-l0-matrix): the optimizer 2x2 with smooth-L0 imp-min (gamma 1->0.01) at 100k steps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Oli's follow-up (task smooth-l0-muon-matrix-100k): swap stock Lp imp-min for SmoothL0ImportanceMinimalityLoss in each matrix cell — coeff 2e-4 (the pile-4L smooth-L0 sweep center, wandb groups smoothl0-impmin-sweep / smoothl0-steplen; numerically same as our Lp cells), gamma 1.0 -> 0.01 linear over the full run (flagship p-594db290 precedent), steps 40k -> 100k. Frequency sub-term kept so the only imp-min delta is the penalty shape. Everything else byte-identical to the parent cell configs. Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- .../smooth_l0_100k/both_muon_sl0_100k.yaml | 160 +++++++++++++++++ .../muon/smooth_l0_100k/ci_muon_sl0_100k.yaml | 162 ++++++++++++++++++ .../muon/smooth_l0_100k/control_sl0_100k.yaml | 161 +++++++++++++++++ .../muon/smooth_l0_100k/muon_1x_sl0_100k.yaml | 159 +++++++++++++++++ 4 files changed, 642 insertions(+) create mode 100644 param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml create mode 100644 param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml create mode 100644 param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml create mode 100644 param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml diff --git a/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml new file mode 100644 index 000000000..ad04b2e22 --- /dev/null +++ b/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml @@ -0,0 +1,160 @@ +# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): both_muon cell +# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L +# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as +# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 +# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. +run_name: jax-pile4l-ppgd-both-muon-sl0-100k +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: 1.0 + gamma_anneal_start_frac: 0.0 + gamma_anneal_final_gamma: 0.01 + gamma_anneal_end_frac: 1.0 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 100000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: smooth-l0-matrix-100k + tags: + - smooth-l0-100k + - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml new file mode 100644 index 000000000..f7a05b59a --- /dev/null +++ b/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml @@ -0,0 +1,162 @@ +# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): ci_muon cell +# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L +# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as +# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 +# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. +run_name: jax-pile4l-ppgd-ci-muon-sl0-100k +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: 1.0 + gamma_anneal_start_frac: 0.0 + gamma_anneal_final_gamma: 0.01 + gamma_anneal_end_frac: 1.0 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 100000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: smooth-l0-matrix-100k + tags: + - smooth-l0-100k + - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml new file mode 100644 index 000000000..aea784dba --- /dev/null +++ b/param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml @@ -0,0 +1,161 @@ +# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): control cell +# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L +# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as +# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 +# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. +run_name: jax-pile4l-ppgd-muon-control-sl0-100k +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: 1.0 + gamma_anneal_start_frac: 0.0 + gamma_anneal_final_gamma: 0.01 + gamma_anneal_end_frac: 1.0 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 100000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: smooth-l0-matrix-100k + tags: + - smooth-l0-100k + - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml new file mode 100644 index 000000000..70d76d835 --- /dev/null +++ b/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml @@ -0,0 +1,159 @@ +# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): muon_1x cell +# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L +# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as +# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 +# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. +run_name: jax-pile4l-ppgd-muon-1x-sl0-100k +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: 1.0 + gamma_anneal_start_frac: 0.0 + gamma_anneal_final_gamma: 0.01 + gamma_anneal_end_frac: 1.0 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 100000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: smooth-l0-matrix-100k + tags: + - smooth-l0-100k + - muon-cifn-matrix From e00d52835e9d2b0381736c7eedb68cbd8310ff25 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sun, 12 Jul 2026 17:30:54 +0000 Subject: [PATCH 09/13] =?UTF-8?q?feat(optim):=20stacked-NS=20muon=20?= =?UTF-8?q?=E2=80=94=20same-shape=20leaves=20batched,=20stack=20axis=20sha?= =?UTF-8?q?rded=20(fast=20muon)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GSPMD lowers per-leaf NS on the ÷N-sharded fp32 masters into per-iteration full-Gram all-reduces with the largest matmul replicated on every device (~24.6 GB serialized fp32 collectives per step for the 4L ci-fn group — the measured 3.3x muon-ci hit). muon_stacked.py batches same-shape muon leaves into one NS with the stack axis sharded over (replicate, fsdp): device-local orthogonalization, one reshard in/out, zero per-iteration collectives (the Kimi parameter-partitioned recipe, GSPMD-native). Config-gated on MuonOptimizerConfig: impl: optax|stacked (default optax = the 07-02/07-11 arms' exact semantics), ns_steps (default 5), ns_dtype (default float32; bfloat16 halves NS compute+comm, stacked-only — masters/momentum stay fp32 per N1). Same MuonState pytree, so checkpoints round-trip across impls. SPEC S20 amended 2026-07-12. build_optimizers now takes the mesh (None for toys/CPU). Tests: stacked-vs-optax update parity on mixed 2D/3D/fallback trees (2 steps, rtol 1e-4); sharded-vs-unsharded parity at 4 sim devices; checkpoint roundtrip + exact-resume with stacked muon on both groups. make check clean. Acceptance configs: both-muon 40k at impl=stacked (fp32) and +bf16-NS, group fast-muon-accept — quality vs p-c01f5833, step time is the metric (0.44s today). Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- param_decomp/SPEC.md | 2 +- param_decomp/configs.py | 20 ++ .../pile4l_ppgd_40k_both_muon_stacked.yaml | 164 +++++++++++++++ ...ile4l_ppgd_40k_both_muon_stacked_bf16.yaml | 166 +++++++++++++++ param_decomp/muon_stacked.py | 193 ++++++++++++++++++ param_decomp/run.py | 2 +- param_decomp/run_state.py | 29 ++- param_decomp/tests/test_checkpoint.py | 50 +++-- param_decomp/tests/test_optim_torch_parity.py | 92 ++++++++- param_decomp_lab/experiments/lm/load_run.py | 2 +- 10 files changed, 696 insertions(+), 24 deletions(-) create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml create mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml create mode 100644 param_decomp/muon_stacked.py diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index 59b1aaa5f..f9f9df71a 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. **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. | +| 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. **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 4b28d4c7c..e4e9553c6 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -603,6 +603,26 @@ class MuonOptimizerConfig(BaseConfig): 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: diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml new file mode 100644 index 000000000..2d731b945 --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml @@ -0,0 +1,164 @@ +# fast-muon acceptance (task muon-sharding-design): both-muon 40k cell with +# stacked NS (fp32) on BOTH optimizer groups. Byte-identical to +# pile4l_ppgd_40k_both_muon.yaml (run p-c01f5833) otherwise — quality curves should +# match it up to reassociation; step time is the acceptance metric (0.44 s today). +run_name: jax-pile4l-ppgd-both-muon-stacked +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + # so disabling it also matches their autotune behavior. + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: fast-muon-accept + tags: + - fast-muon-accept + - muon-cifn-matrix diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml new file mode 100644 index 000000000..898feed63 --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml @@ -0,0 +1,166 @@ +# fast-muon acceptance (task muon-sharding-design): both-muon 40k cell with +# stacked NS + bf16 orthogonalization on BOTH optimizer groups. Byte-identical to +# pile4l_ppgd_40k_both_muon.yaml (run p-c01f5833) otherwise — quality curves should +# match it up to reassociation; step time is the acceptance metric (0.44 s today). +run_name: jax-pile4l-ppgd-both-muon-stacked-bf16 +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + ns_dtype: bfloat16 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + ns_dtype: bfloat16 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + # so disabling it also matches their autotune behavior. + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: fast-muon-accept + tags: + - fast-muon-accept + - muon-cifn-matrix diff --git a/param_decomp/muon_stacked.py b/param_decomp/muon_stacked.py new file mode 100644 index 000000000..63c7572b5 --- /dev/null +++ b/param_decomp/muon_stacked.py @@ -0,0 +1,193 @@ +"""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] + grouped = jnp.concatenate(stacks, axis=0) + 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.astype(ns_dtype), + jnp.asarray(ns_coeffs, ns_dtype), + ns_steps=ns_steps, + dimension_numbers=_NS_DIMS, + ).astype(grouped.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 bd219744a..31b9cb63a 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -465,7 +465,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 7e93ad24a..c62b4e43d 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -29,6 +29,7 @@ PDConfig, ) from param_decomp.lm import DecomposedModel +from param_decomp.muon_stacked import stacked_muon from param_decomp.recon import ( PersistentSources, build_loss_terms, @@ -100,32 +101,48 @@ 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.""" + 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(): + 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 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). @@ -137,13 +154,15 @@ def build_optimizers(pd: PDConfig): pd.components_optimizer.lr_schedule.start_val, pd.steps, alpha=0.1 ) sched_ci = torch_cosine_schedule(pd.ci_fn_optimizer.lr_schedule.start_val, pd.steps, alpha=0.1) - opt_vu = _optimizer_with_clip(pd.components_optimizer, sched_vu, muon_dimension_numbers=None) + 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) + 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 6aea1963b..23cb235a7 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,6 +39,7 @@ 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 @@ -95,25 +97,34 @@ def _chunkwise_arch(lm: DecomposedModel, cfg: LlamaConfig) -> ChunkwiseTransform ) -def _build(seed: int, muon_components: bool = False, muon_ci_fn: bool = False): +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)) - inner_vu = ( - optax.contrib.muon(1e-3, consistent_rms=0.2) - if muon_components - else 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 = ( - optax.contrib.muon( - 1e-3, - consistent_rms=0.2, - muon_weight_dimension_numbers=chunk_stacked_muon_dimension_numbers, - ) + muon_impl(chunk_stacked_muon_dimension_numbers) if muon_ci_fn else optax.adamw(1e-3, weight_decay=0.0) ) @@ -156,9 +167,11 @@ def _build(seed: int, muon_components: bool = False, muon_ci_fn: bool = False): def _roundtrip_and_exact_resume( - tmp_path: Path, muon_components: bool, muon_ci_fn: bool = False + 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) + 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)) @@ -166,7 +179,9 @@ def _roundtrip_and_exact_resume( save_state(mgr, 2, state) # Restore onto a DIFFERENTLY-seeded reference: every leaf must come from disk. - _, fresh, _, _ = _build(seed=7, muon_components=muon_components, muon_ci_fn=muon_ci_fn) + _, 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 @@ -201,6 +216,13 @@ def test_muon_ci_fn_roundtrip_and_exact_resume(tmp_path: Path): _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 b8e90c158..b4f11c532 100644 --- a/param_decomp/tests/test_optim_torch_parity.py +++ b/param_decomp/tests/test_optim_torch_parity.py @@ -114,7 +114,7 @@ def test_muon_orthogonalizes_2d_leaves_and_adam_falls_back_elsewhere(): grad_clip_norm=0.01, ) lr = 1e-3 - opt = _optimizer_with_clip(muon_cfg, lambda count: jnp.float32(lr), None) + 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 = { @@ -155,7 +155,7 @@ def test_muon_chunk_stacked_dimension_numbers_orthogonalize_3d_and_adam_2d_bias_ ) lr = 1e-3 opt = _optimizer_with_clip( - muon_cfg, lambda count: jnp.float32(lr), chunk_stacked_muon_dimension_numbers + muon_cfg, lambda count: jnp.float32(lr), chunk_stacked_muon_dimension_numbers, mesh=None ) key = jax.random.key(0) n_chunks = 3 @@ -184,6 +184,94 @@ def test_muon_chunk_stacked_dimension_numbers_orthogonalize_3d_and_adam_2d_bias_ ) +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) diff --git a/param_decomp_lab/experiments/lm/load_run.py b/param_decomp_lab/experiments/lm/load_run.py index 8ee3cfa2a..f19222125 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 From d8e4571614e0286f5de4fe2d979da61c451966b2 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sun, 12 Jul 2026 17:47:58 +0000 Subject: [PATCH 10/13] perf(muon-stacked): cast to ns_dtype before the stack-axis reshard The ingress reshard was moving fp32 bytes even with bf16 NS (cast sat after the sharding constraint); casting first halves the ingress bytes for ns_dtype=bfloat16. fp32 NS unchanged (cast is a no-op). Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- param_decomp/muon_stacked.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/param_decomp/muon_stacked.py b/param_decomp/muon_stacked.py index 63c7572b5..f31a37912 100644 --- a/param_decomp/muon_stacked.py +++ b/param_decomp/muon_stacked.py @@ -75,16 +75,19 @@ def _grouped_newton_schulz( for shape, indices in groups.items(): stacks = [canon[i][0] for i in indices] sizes = [s.shape[0] for s in stacks] - grouped = jnp.concatenate(stacks, axis=0) + 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.astype(ns_dtype), + grouped, jnp.asarray(ns_coeffs, ns_dtype), ns_steps=ns_steps, dimension_numbers=_NS_DIMS, - ).astype(grouped.dtype) + ).astype(out_dtype) offset = 0 for i, size in zip(indices, sizes, strict=True): piece = orthogonalized[offset : offset + size] From dbe66425df7feee7ef17171396a09260531377ed Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sun, 12 Jul 2026 19:50:02 +0000 Subject: [PATCH 11/13] config(smooth-l0-matrix): stacked-bf16 relaunch variants for the three muon cells The original sl0-100k muon arms launched on the per-leaf optax impl (0.40-0.44 s/step for the ci cells); fresh 100k relaunches on stacked-bf16 (~0.23 s/step) finish ~2-3h sooner than the originals' remaining tail, and give all three muon cells identical NS semantics (the flagship-intended mode). Gated on the fast-muon acceptance arms' quality-parity check vs p-c01f5833. adam/adam has no muon group and rides on. Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- .../both_muon_sl0_100k_fast.yaml | 162 ++++++++++++++++++ .../smooth_l0_100k/ci_muon_sl0_100k_fast.yaml | 162 ++++++++++++++++++ .../smooth_l0_100k/muon_1x_sl0_100k_fast.yaml | 159 +++++++++++++++++ 3 files changed, 483 insertions(+) create mode 100644 param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml create mode 100644 param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml create mode 100644 param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml diff --git a/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml b/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml new file mode 100644 index 000000000..060e543d9 --- /dev/null +++ b/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml @@ -0,0 +1,162 @@ +# both_muon smooth-L0 100k cell RELAUNCHED on stacked-bf16 NS (impl: stacked, +# ns_dtype: bfloat16) — supersedes the old-impl arm for wall-clock (~6.5h vs ~9h +# remaining); only the muon NS impl/dtype differs from both_muon_sl0_100k.yaml. +run_name: jax-pile4l-ppgd-both-muon-sl0-100k-fast +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + ns_dtype: bfloat16 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + ns_dtype: bfloat16 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: 1.0 + gamma_anneal_start_frac: 0.0 + gamma_anneal_final_gamma: 0.01 + gamma_anneal_end_frac: 1.0 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 100000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: smooth-l0-matrix-100k + tags: + - smooth-l0-100k + - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml b/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml new file mode 100644 index 000000000..81f2d4153 --- /dev/null +++ b/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml @@ -0,0 +1,162 @@ +# ci_muon smooth-L0 100k cell RELAUNCHED on stacked-bf16 NS (impl: stacked, +# ns_dtype: bfloat16) — supersedes the old-impl arm for wall-clock (~6.5h vs ~9h +# remaining); only the muon NS impl/dtype differs from ci_muon_sl0_100k.yaml. +run_name: jax-pile4l-ppgd-ci-muon-sl0-100k-fast +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + ns_dtype: bfloat16 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: 1.0 + gamma_anneal_start_frac: 0.0 + gamma_anneal_final_gamma: 0.01 + gamma_anneal_end_frac: 1.0 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 100000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main + # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: smooth-l0-matrix-100k + tags: + - smooth-l0-100k + - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml b/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml new file mode 100644 index 000000000..e0d6959d0 --- /dev/null +++ b/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml @@ -0,0 +1,159 @@ +# muon_1x smooth-L0 100k cell RELAUNCHED on stacked-bf16 NS (impl: stacked, +# ns_dtype: bfloat16) — supersedes the old-impl arm for wall-clock (~6.5h vs ~9h +# remaining); only the muon NS impl/dtype differs from muon_1x_sl0_100k.yaml. +run_name: jax-pile4l-ppgd-muon-1x-sl0-100k-fast +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + type: muon + consistent_rms: 0.2 + impl: stacked + ns_dtype: bfloat16 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: 1.0 + gamma_anneal_start_frac: 0.0 + gamma_anneal_final_gamma: 0.01 + gamma_anneal_end_frac: 1.0 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 100000 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: smooth-l0-matrix-100k + tags: + - smooth-l0-100k + - muon-cifn-matrix From 9d41dd8f1c37056b09b3961c8db42d1591f8f2ba Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Mon, 13 Jul 2026 19:46:13 +0000 Subject: [PATCH 12/13] config(probe): 400-step autotune-cache step-time probe pair Replicates the +0.08s/step regression seen in every post-07-11 launch: two adamw/adamw arms identical except JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES=none. HLO dumps land in /hlo for diffing. Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- .../configs/muon/probe_autotune_default.yaml | 157 +++++++++++++++++ .../configs/muon/probe_autotune_nocache.yaml | 160 ++++++++++++++++++ 2 files changed, 317 insertions(+) create mode 100644 param_decomp/configs/muon/probe_autotune_default.yaml create mode 100644 param_decomp/configs/muon/probe_autotune_nocache.yaml diff --git a/param_decomp/configs/muon/probe_autotune_default.yaml b/param_decomp/configs/muon/probe_autotune_default.yaml new file mode 100644 index 000000000..9cf7d9940 --- /dev/null +++ b/param_decomp/configs/muon/probe_autotune_default.yaml @@ -0,0 +1,157 @@ +# autotune-cache step-time probe (default): 400-step adamw/adamw replication of the +# +0.08s/step regression seen in every post-07-11 launch. Identical to +# pile4l_ppgd_40k_control.yaml except steps/logging. +# HLO lands in /hlo (trainer dumps by default) for diffing. +run_name: pile4l-ppgd_uv-adam_ci-adam_probe-default +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 100 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: false +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400 +runtime: + remat_recon_forwards: false + dp: 8 +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: autotune-probe + tags: + - hlo-probe diff --git a/param_decomp/configs/muon/probe_autotune_nocache.yaml b/param_decomp/configs/muon/probe_autotune_nocache.yaml new file mode 100644 index 000000000..621991be5 --- /dev/null +++ b/param_decomp/configs/muon/probe_autotune_nocache.yaml @@ -0,0 +1,160 @@ +# autotune-cache step-time probe (nocache): 400-step adamw/adamw replication of the +# +0.08s/step regression seen in every post-07-11 launch. Identical to +# pile4l_ppgd_40k_control.yaml except steps/logging and the =none env var. +# HLO lands in /hlo (trainer dumps by default) for diffing. +run_name: pile4l-ppgd_uv-adam_ci-adam_probe-nocache +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 100 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 1000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: shared_across_batch + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: false +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + eps: 1.0e-12 + p_anneal_end_frac: 1.0 + p_anneal_final_p: 0.4 + p_anneal_start_frac: 0.0 + pnorm: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_samples: 1 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400 +runtime: + remat_recon_forwards: false + dp: 8 + launch_env: + env: + JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none +target: + output_extract: 0 + weights_dtype: bfloat16 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 +wandb: + entity: null + project: param-decomp + group: autotune-probe + tags: + - hlo-probe From c0727a8080731e6a4e6c0f24898b303854ca381f Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Mon, 13 Jul 2026 20:56:17 +0000 Subject: [PATCH 13/13] config: drop the muon experiment yamls per the config policy (CONFIGS.md / #983) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sweep one-offs don't get committed — provenance for the launched runs lives in the run-dir pinned launch_config.yaml, the refs/runs/snapshot/* refs, and wandb. No canonical muon seat yet: muon stays config-gated experimental until it's promoted to a default recipe, at which point a seat row lands in CONFIGS.md. The experiment configs remain on the bridge/task-* experiment branches for history. Crew-Address: slack/C08T7UV4449/1783730484.936959 Co-Authored-By: Claude Fable 5 --- .../muon/impmin_sweep/both_muon_imp05x.yaml | 159 ----------------- .../muon/impmin_sweep/both_muon_imp2x.yaml | 159 ----------------- .../muon/impmin_sweep/ci_muon_imp05x.yaml | 161 ----------------- .../muon/impmin_sweep/ci_muon_imp2x.yaml | 161 ----------------- .../muon/impmin_sweep/control_imp05x.yaml | 160 ----------------- .../muon/impmin_sweep/control_imp2x.yaml | 160 ----------------- .../muon/impmin_sweep/muon_1x_imp05x.yaml | 158 ----------------- .../muon/impmin_sweep/muon_1x_imp2x.yaml | 158 ----------------- .../muon/pile4l_ppgd_40k_adamw_4x.yaml | 157 ----------------- .../muon/pile4l_ppgd_40k_both_muon.yaml | 161 ----------------- .../pile4l_ppgd_40k_both_muon_stacked.yaml | 164 ----------------- ...ile4l_ppgd_40k_both_muon_stacked_bf16.yaml | 166 ------------------ .../configs/muon/pile4l_ppgd_40k_ci_muon.yaml | 164 ----------------- .../configs/muon/pile4l_ppgd_40k_control.yaml | 157 ----------------- .../configs/muon/pile4l_ppgd_40k_muon_1x.yaml | 157 ----------------- .../configs/muon/pile4l_ppgd_40k_muon_4x.yaml | 157 ----------------- .../configs/muon/probe_autotune_default.yaml | 157 ----------------- .../configs/muon/probe_autotune_nocache.yaml | 160 ----------------- .../smooth_l0_100k/both_muon_sl0_100k.yaml | 160 ----------------- .../both_muon_sl0_100k_fast.yaml | 162 ----------------- .../muon/smooth_l0_100k/ci_muon_sl0_100k.yaml | 162 ----------------- .../smooth_l0_100k/ci_muon_sl0_100k_fast.yaml | 162 ----------------- .../muon/smooth_l0_100k/control_sl0_100k.yaml | 161 ----------------- .../muon/smooth_l0_100k/muon_1x_sl0_100k.yaml | 159 ----------------- .../smooth_l0_100k/muon_1x_sl0_100k_fast.yaml | 159 ----------------- 25 files changed, 4001 deletions(-) delete mode 100644 param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml delete mode 100644 param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml delete mode 100644 param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml delete mode 100644 param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml delete mode 100644 param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml delete mode 100644 param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml delete mode 100644 param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml delete mode 100644 param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_control.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml delete mode 100644 param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml delete mode 100644 param_decomp/configs/muon/probe_autotune_default.yaml delete mode 100644 param_decomp/configs/muon/probe_autotune_nocache.yaml delete mode 100644 param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml delete mode 100644 param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml delete mode 100644 param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml delete mode 100644 param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml delete mode 100644 param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml delete mode 100644 param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml delete mode 100644 param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml diff --git a/param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml deleted file mode 100644 index 4a0c5582c..000000000 --- a/param_decomp/configs/muon/impmin_sweep/both_muon_imp05x.yaml +++ /dev/null @@ -1,159 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): both_muon cell at imp-min -# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_both_muon.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-both-muon-imp05x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0001 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml deleted file mode 100644 index 85a28387f..000000000 --- a/param_decomp/configs/muon/impmin_sweep/both_muon_imp2x.yaml +++ /dev/null @@ -1,159 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): both_muon cell at imp-min -# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_both_muon.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-both-muon-imp2x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0004 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml deleted file mode 100644 index 284ffaaad..000000000 --- a/param_decomp/configs/muon/impmin_sweep/ci_muon_imp05x.yaml +++ /dev/null @@ -1,161 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): ci_muon cell at imp-min -# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_ci_muon.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-ci-muon-imp05x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0001 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml deleted file mode 100644 index 0e26d3e3c..000000000 --- a/param_decomp/configs/muon/impmin_sweep/ci_muon_imp2x.yaml +++ /dev/null @@ -1,161 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): ci_muon cell at imp-min -# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_ci_muon.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-ci-muon-imp2x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0004 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml deleted file mode 100644 index ac97e547b..000000000 --- a/param_decomp/configs/muon/impmin_sweep/control_imp05x.yaml +++ /dev/null @@ -1,160 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): control cell at imp-min -# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_control.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-muon-control-imp05x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0001 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml deleted file mode 100644 index 955f81515..000000000 --- a/param_decomp/configs/muon/impmin_sweep/control_imp2x.yaml +++ /dev/null @@ -1,160 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): control cell at imp-min -# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_control.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-muon-control-imp2x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0004 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml b/param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml deleted file mode 100644 index fd0484fd7..000000000 --- a/param_decomp/configs/muon/impmin_sweep/muon_1x_imp05x.yaml +++ /dev/null @@ -1,158 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): muon_1x cell at imp-min -# coeff 0.0001 (imp05x); byte-derived from pile4l_ppgd_40k_muon_1x.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-muon-1x-imp05x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0001 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml b/param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml deleted file mode 100644 index 827cdef01..000000000 --- a/param_decomp/configs/muon/impmin_sweep/muon_1x_imp2x.yaml +++ /dev/null @@ -1,158 +0,0 @@ -# impmin sweep (bridge/task muon-cifn-matrix follow-up): muon_1x cell at imp-min -# coeff 0.0004 (imp2x); byte-derived from pile4l_ppgd_40k_muon_1x.yaml, only run_name / coeff / -# (env fix + tags for the try-muon-era bases) changed. -run_name: jax-pile4l-ppgd-muon-1x-imp2x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0004 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix - - impmin-sweep diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml deleted file mode 100644 index 29d76bab3..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_adamw_4x.yaml +++ /dev/null @@ -1,157 +0,0 @@ -# try-muon A/B: adamw at 4x lr (2e-4) — the LR-vs-optimizer disambiguation control for muon_4x. -# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this -# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix -# cherry-picked here). Control arm: canonical AdamW, trajectory-identical to stock. -run_name: jax-pile4l-ppgd-adamw-4x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 2.0e-04 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: try-muon - tags: - - try-muon diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml deleted file mode 100644 index f053d3958..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml +++ /dev/null @@ -1,161 +0,0 @@ -# muon x {UV, ci-fn} matrix (bridge/task-muon-cifn-matrix): byte-derived from the try-muon -# muon-1x arm p-cedbbf7c's config (= pile_ppgd_bsc.yaml, steps 400k->40k, dp 16->8, seed 0, -# components_optimizer muon @ lr 5e-5 consistent_rms 0.2), with ONLY ci_fn_optimizer ALSO -# flipped to muon (same settings, still unclipped per S19). Matrix cell: UV=muon, ci-fn=muon. -run_name: jax-pile4l-ppgd-both-muon -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - # so disabling it also matches their autotune behavior. - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml deleted file mode 100644 index c02aab8e9..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked.yaml +++ /dev/null @@ -1,164 +0,0 @@ -# fast-muon acceptance (task muon-sharding-design): both-muon 40k cell with -# stacked NS (fp32) on BOTH optimizer groups. Byte-identical to -# pile4l_ppgd_40k_both_muon.yaml (run p-c01f5833) otherwise — quality curves should -# match it up to reassociation; step time is the acceptance metric (0.44 s today). -run_name: jax-pile4l-ppgd-both-muon-stacked -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - # so disabling it also matches their autotune behavior. - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: fast-muon-accept - tags: - - fast-muon-accept - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml deleted file mode 100644 index 00b0f8a53..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon_stacked_bf16.yaml +++ /dev/null @@ -1,166 +0,0 @@ -# fast-muon acceptance (task muon-sharding-design): both-muon 40k cell with -# stacked NS + bf16 orthogonalization on BOTH optimizer groups. Byte-identical to -# pile4l_ppgd_40k_both_muon.yaml (run p-c01f5833) otherwise — quality curves should -# match it up to reassociation; step time is the acceptance metric (0.44 s today). -run_name: jax-pile4l-ppgd-both-muon-stacked-bf16 -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - ns_dtype: bfloat16 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - ns_dtype: bfloat16 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - # so disabling it also matches their autotune behavior. - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: fast-muon-accept - tags: - - fast-muon-accept - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml deleted file mode 100644 index 8eebd1166..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml +++ /dev/null @@ -1,164 +0,0 @@ -# muon x {UV, ci-fn} matrix (bridge/task-muon-cifn-matrix): byte-derived from the try-muon -# control arm p-d0f66d3b's config (= pile_ppgd_bsc.yaml, steps 400k->40k, dp 16->8, seed 0), -# with ONLY ci_fn_optimizer flipped to muon (consistent_rms 0.2, same lr/schedule, still -# unclipped per S19). Matrix cell: UV=adamw, ci-fn=muon. Chunk-stacked 3D ci-fn leaves are -# muon'd via chunk_stacked_muon_dimension_numbers; 2D bias stacks take the Adam fallback. -run_name: jax-pile4l-ppgd-ci-muon -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # The shared xla_gpu_per_fusion_autotune_cache_dir was created group-unwritable - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - # so disabling it also matches their autotune behavior. - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: muon-cifn-matrix - tags: - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_control.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_control.yaml deleted file mode 100644 index 2ea2740f8..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_control.yaml +++ /dev/null @@ -1,157 +0,0 @@ -# try-muon A/B (bridge/task-try-muon): byte-derived from baseline p-5d480c2d's pinned -# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this -# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix -# cherry-picked here). Control arm: canonical AdamW, trajectory-identical to stock. -run_name: jax-pile4l-ppgd-muon-control -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: try-muon - tags: - - try-muon diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml deleted file mode 100644 index b6ee1504f..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_muon_1x.yaml +++ /dev/null @@ -1,157 +0,0 @@ -# try-muon A/B (bridge/task-try-muon): byte-derived from baseline p-5d480c2d's pinned -# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this -# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix -# cherry-picked here). Muon arm (1x = lr 5.0e-05): components_optimizer type=muon, consistent_rms 0.2 (update -# RMS matched to AdamW's per the Kimi recipe, so lr is on the AdamW scale), same cosine -# schedule + S19 clip; ci_fn_optimizer stays canonical AdamW for clean attribution. -run_name: jax-pile4l-ppgd-muon-1x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: try-muon - tags: - - try-muon diff --git a/param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml deleted file mode 100644 index 9a357ff5f..000000000 --- a/param_decomp/configs/muon/pile4l_ppgd_40k_muon_4x.yaml +++ /dev/null @@ -1,157 +0,0 @@ -# try-muon A/B (bridge/task-try-muon): byte-derived from baseline p-5d480c2d's pinned -# launch_config (= pile_ppgd_bsc.yaml with steps 400k->40k, dp 16->8), relaunched from this -# branch because p-5d480c2d died at step 1000 on the hidden-acts-eval crash (#919, fix -# cherry-picked here). Muon arm (4x = lr 2.0e-04): components_optimizer type=muon, consistent_rms 0.2 (update -# RMS matched to AdamW's per the Kimi recipe, so lr is on the AdamW scale), same cosine -# schedule + S19 clip; ci_fn_optimizer stays canonical AdamW for clean attribution. -run_name: jax-pile4l-ppgd-muon-4x -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 2.0e-04 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40000 -runtime: - remat_recon_forwards: false - dp: 8 -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: try-muon - tags: - - try-muon diff --git a/param_decomp/configs/muon/probe_autotune_default.yaml b/param_decomp/configs/muon/probe_autotune_default.yaml deleted file mode 100644 index f6b592420..000000000 --- a/param_decomp/configs/muon/probe_autotune_default.yaml +++ /dev/null @@ -1,157 +0,0 @@ -# autotune-cache step-time probe (default): 400-step adamw/adamw replication of the -# +0.08s/step regression seen in every post-07-11 launch. Identical to -# pile4l_ppgd_40k_control.yaml except steps/logging. -# HLO lands in /hlo (trainer dumps by default) for diffing. -run_name: pile4l-ppgd_uv-adam_ci-adam_probe-default -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 100 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: false -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 400 -runtime: - remat_recon_forwards: false - dp: 8 -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: autotune-probe - tags: - - hlo-probe diff --git a/param_decomp/configs/muon/probe_autotune_nocache.yaml b/param_decomp/configs/muon/probe_autotune_nocache.yaml deleted file mode 100644 index 454a88fd2..000000000 --- a/param_decomp/configs/muon/probe_autotune_nocache.yaml +++ /dev/null @@ -1,160 +0,0 @@ -# autotune-cache step-time probe (nocache): 400-step adamw/adamw replication of the -# +0.08s/step regression seen in every post-07-11 launch. Identical to -# pile4l_ppgd_40k_control.yaml except steps/logging and the =none env var. -# HLO lands in /hlo (trainer dumps by default) for diffing. -run_name: pile4l-ppgd_uv-adam_ci-adam_probe-nocache -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 100 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: false -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - eps: 1.0e-12 - pnorm: - start_val: 2.0 - fn_type: linear - final_val_frac: 0.2 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 400 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: autotune-probe - tags: - - hlo-probe diff --git a/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml deleted file mode 100644 index 74f107b11..000000000 --- a/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k.yaml +++ /dev/null @@ -1,160 +0,0 @@ -# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): both_muon cell -# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L -# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as -# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 -# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. -run_name: jax-pile4l-ppgd-both-muon-sl0-100k -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - gamma: - start_val: 1.0 - fn_type: linear - final_val_frac: 0.01 - type: SmoothL0ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 100000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: smooth-l0-matrix-100k - tags: - - smooth-l0-100k - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml b/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml deleted file mode 100644 index d341974c1..000000000 --- a/param_decomp/configs/muon/smooth_l0_100k/both_muon_sl0_100k_fast.yaml +++ /dev/null @@ -1,162 +0,0 @@ -# both_muon smooth-L0 100k cell RELAUNCHED on stacked-bf16 NS (impl: stacked, -# ns_dtype: bfloat16) — supersedes the old-impl arm for wall-clock (~6.5h vs ~9h -# remaining); only the muon NS impl/dtype differs from both_muon_sl0_100k.yaml. -run_name: jax-pile4l-ppgd-both-muon-sl0-100k-fast -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - ns_dtype: bfloat16 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - ns_dtype: bfloat16 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - gamma: - start_val: 1.0 - fn_type: linear - final_val_frac: 0.01 - type: SmoothL0ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 100000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: smooth-l0-matrix-100k - tags: - - smooth-l0-100k - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml deleted file mode 100644 index 36eb7f103..000000000 --- a/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k.yaml +++ /dev/null @@ -1,162 +0,0 @@ -# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): ci_muon cell -# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L -# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as -# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 -# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. -run_name: jax-pile4l-ppgd-ci-muon-sl0-100k -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - gamma: - start_val: 1.0 - fn_type: linear - final_val_frac: 0.01 - type: SmoothL0ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 100000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: smooth-l0-matrix-100k - tags: - - smooth-l0-100k - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml b/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml deleted file mode 100644 index 7e0d02992..000000000 --- a/param_decomp/configs/muon/smooth_l0_100k/ci_muon_sl0_100k_fast.yaml +++ /dev/null @@ -1,162 +0,0 @@ -# ci_muon smooth-L0 100k cell RELAUNCHED on stacked-bf16 NS (impl: stacked, -# ns_dtype: bfloat16) — supersedes the old-impl arm for wall-clock (~6.5h vs ~9h -# remaining); only the muon NS impl/dtype differs from ci_muon_sl0_100k.yaml. -run_name: jax-pile4l-ppgd-ci-muon-sl0-100k-fast -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - ns_dtype: bfloat16 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - gamma: - start_val: 1.0 - fn_type: linear - final_val_frac: 0.01 - type: SmoothL0ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 100000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - # (2026-07-11, PERMISSION_DENIED at first step for any non-owner run); the main - # executable cache is unaffected. The 2026-07-03 sibling arms predate that dir, - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: smooth-l0-matrix-100k - tags: - - smooth-l0-100k - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml deleted file mode 100644 index 54cd1b4c5..000000000 --- a/param_decomp/configs/muon/smooth_l0_100k/control_sl0_100k.yaml +++ /dev/null @@ -1,161 +0,0 @@ -# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): control cell -# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L -# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as -# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 -# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. -run_name: jax-pile4l-ppgd-muon-control-sl0-100k -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - gamma: - start_val: 1.0 - fn_type: linear - final_val_frac: 0.01 - type: SmoothL0ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 100000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: smooth-l0-matrix-100k - tags: - - smooth-l0-100k - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml b/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml deleted file mode 100644 index da3772525..000000000 --- a/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k.yaml +++ /dev/null @@ -1,159 +0,0 @@ -# smooth-L0 x optimizer matrix at 100k (task smooth-l0-muon-matrix-100k): muon_1x cell -# with SmoothL0ImportanceMinimalityLoss replacing stock Lp — coeff 2e-4 (the pile-4L -# smooth-L0 sweep center, groups smoothl0-impmin-sweep/smoothl0-steplen; same value as -# our Lp cells), gamma 1.0 -> 0.01 linear full-window (Oli's ask; flagship p-594db290 -# precedent). steps 40k -> 100k. Everything else byte-identical to the parent cell. -run_name: jax-pile4l-ppgd-muon-1x-sl0-100k -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - gamma: - start_val: 1.0 - fn_type: linear - final_val_frac: 0.01 - type: SmoothL0ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 100000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: smooth-l0-matrix-100k - tags: - - smooth-l0-100k - - muon-cifn-matrix diff --git a/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml b/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml deleted file mode 100644 index eaf15e831..000000000 --- a/param_decomp/configs/muon/smooth_l0_100k/muon_1x_sl0_100k_fast.yaml +++ /dev/null @@ -1,159 +0,0 @@ -# muon_1x smooth-L0 100k cell RELAUNCHED on stacked-bf16 NS (impl: stacked, -# ns_dtype: bfloat16) — supersedes the old-impl arm for wall-clock (~6.5h vs ~9h -# remaining); only the muon NS impl/dtype differs from muon_1x_sl0_100k.yaml. -run_name: jax-pile4l-ppgd-muon-1x-sl0-100k-fast -cadence: - keep_last_n_checkpoints: 2 - save_every: 5000 - train_log_every: 200 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 1000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: shared_across_batch - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 10000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - type: chunkwise_transformer - blocks_per_chunk: 1 - d_model: 2048 - n_blocks: 8 - n_heads: 16 - mlp_hidden: 8192 - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - type: muon - consistent_rms: 0.2 - impl: stacked - ns_dtype: bfloat16 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 400 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - frequency: - coeff: 0.0001 - reference_token_count: 65536 - coeff: 0.0002 - gamma: - start_val: 1.0 - fn_type: linear - final_val_frac: 0.01 - type: SmoothL0ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_samples: 1 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: per_batch_per_position - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 100000 -runtime: - remat_recon_forwards: false - dp: 8 - launch_env: - env: - JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: none -target: - output_extract: 0 - weights_dtype: bfloat16 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 -wandb: - entity: null - project: param-decomp - group: smooth-l0-matrix-100k - tags: - - smooth-l0-100k - - muon-cifn-matrix