From e00d52835e9d2b0381736c7eedb68cbd8310ff25 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Sun, 12 Jul 2026 17:30:54 +0000 Subject: [PATCH 1/4] =?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 2/4] 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 3/4] 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 4/4] 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