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/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 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..44127f39e --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_both_muon.yaml @@ -0,0 +1,161 @@ +# 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 + 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_ci_muon.yaml b/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml new file mode 100644 index 000000000..bcaf0f19b --- /dev/null +++ b/param_decomp/configs/muon/pile4l_ppgd_40k_ci_muon.yaml @@ -0,0 +1,164 @@ +# 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 + 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/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 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)