diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index 23d0844c9..40c3b8f3a 100644 --- a/param_decomp/SPEC.md +++ b/param_decomp/SPEC.md @@ -155,7 +155,7 @@ def imp_min_terms(ci_upper, pnorm, reference_token_count): # per-site grouping def stochastic_recon_loss(components, ci_lower, residual, clean_output): total, n_forwards = 0, 0 for (live_sites, SAMPLE_ROUTING) in RECON_PLAN: - for routes in SAMPLE_ROUTING(key, [B,T]): # fresh per step (R1,S11) + for routes in SAMPLE_ROUTING(key, [B,T], step): # fresh per step (R1,S11) masks, delta_masks = make_masks(ci_lower_s, source_s ~ U[0,1]^[B,T,C+1]) ∀ s ∈ live_sites total += kl_per_position(masked_forward(residual, live_sites, masks, delta_masks, routes), clean_output) @@ -280,7 +280,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g | S8' | Imp-min contributes `imp_coeff·lp + freq_coeff·freq` with INDEPENDENT coefficients; the frequency normalizer `a' = reference_token_count` is explicit (batch-invariant at fixed firing rate), not the implicit `B·T`. `freq` is absent when no `frequency` is configured. `a' = B·T` recovers the old rolled `imp_coeff·Σ_c f·(1 + beta·log2(1 + B·T·f))` with `freq_coeff = imp_coeff·beta`. | | S9 | `pnorm(step)` follows its `ScheduleConfig` (canonical: linear `2.0 → 0.4` over the run, `final_val_frac=0.2`), evaluated by `schedule.scheduled_value_traced`; `eps` sits inside the power. The smooth-L0 `gamma(step)` (S9′) follows the same rule. Constant-`p` is `fn_type=constant`; warmup on `p`/`gamma` is refused (`build_loss_terms`). **AMENDED 2026-07-02** (#915): the former windowed linear anneal (`{p,gamma}_anneal_{start,end}_frac`) is retired — every config on record used the trivial `[0, 1]` window — and the anneal now shares the schedule's `decay_steps − 1` denominator (S20), reaching the final value AT `step = steps − 1` instead of one step past the run (a ≤O(1/steps) per-step trajectory change vs the old `step / steps`). | | S10′ | The recon objective is a static tuple of coefficiented loss TERMS (one per configured recon loss metric, in config order). Each term is a static plan of `(live_sites, SAMPLE_ROUTING, MASK_SOURCE)` entries; the term's loss = mean over ALL its forwards (every draw of every entry) of `kl_per_position`; the total adds `coeff · term` per term. Plan structures (live-sets, sampler identities, family sizes, strategy kinds) are fixed across steps. The §4 pseudocode shows the production two-term instantiation (`stochastic_recon_loss` + `adversarial_recon_loss`). Recon KL direction is pinned by S25; the mean-over-forwards ≡ accumulator identity by S26. | -| S11 | `uniform_k_routing`, per position: `k ~ U{1..|live_sites|}` then a uniform `k`-subset of the live sites routes True; non-live sites are not live at all. Routing draws are fresh per step, sampled inside the step. | +| S11 | `uniform_k_routing`, per position: `k ~ U{1..|live_sites|}` then a uniform `k`-subset of the live sites routes True; non-live sites are not live at all. Routing draws are fresh per step, sampled inside the step. **AMENDED 2026-07-12** (additive): three more routing samplers share the rule — `fixed_k_routing` (per position, a uniform subset of exactly `k` live sites, `1 ≤ k ≤ |live_sites|`), `static_probability_routing` (per position × live site, independent `Bernoulli(p)` — the torch `StaticProbabilityRouter`), and `scheduled_probability_routing` (`Bernoulli(p(step))`, `p` a `ScheduleConfig` evaluated by `scheduled_value_traced`; both schedule endpoints must be probabilities, an increasing ramp uses `final_val_frac > 1`). SAMPLE_ROUTING therefore takes the traced step: `(key, [B,T], step_f32)`; only scheduled samplers read it. | | S12′ | An adversarial term's loss forward consumes its sources as LEAVES (no ascent-graph history); gradient flows to components and (through `ci_lower`) to the CI fn — and, for persistent sources, to the leaves themselves (S14′). The PRODUCTION adversarial term masks ALL sites and routes everywhere; subset-routed adversarial terms route per their plan. | | S13′ | Per persistent term: source updates per training step = `n_warmup + 1`, all through THAT term's persistent SRC_STEP optimizer state; its source LR schedule advances once per training step. The source LR is `scheduled_value_traced` over the term's `ScheduleConfig` (constant-after-warmup only — `build_loss_terms` and the lab conversion both refuse decay). **AMENDED 2026-07-02** (#915): the former `warmup_pct==0` accepted seam (JAX clamped `warmup_steps = max(floor(...), 1)` → source LR `=0` at step 0) is retired — at `warmup_pct==0` JAX now matches torch's full LR at step 0. Production uses 2.5% warmup and is unaffected either way. | | S14′ | Each persistent term's final ascent gradient comes from the SAME graph as the main backward (pre-update components, live `ci_lower`), unscaled by THAT term's coeff. It is applied after backward; it must not use post-update params. | @@ -310,7 +310,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g |---|---|---| | `RECON_TERMS` | any static tuple of coefficiented terms, each a static plan of `(live_sites, SAMPLE_ROUTING, MASK_SOURCE)` entries: `subset_chunk_plan` · `per_site_plan` (the torch "layerwise" shape) · `all_sites_plan` · custom subset families (pairs, covers, …); built from the shared torch loss configs by `build_loss_terms` | ★ stochastic subset term + persistent-PGD term | | `MASK_SOURCE` | `stochastic` (fresh U[0,1]/Bernoulli per draw) · `constant(v)` (`v=0` CI-masked, `v=1` unmasked; no delta path) · `fresh_pgd(init, n_steps, step_size, scope)` · `persistent(state_key)` | ★ stochastic + persistent | -| `SAMPLE_ROUTING` | `(key, [B,T]) → tuple of routing draws`, statically sized; draws may be jointly sampled (independent repeats, antithetic/complementary subsets, per-step covers) | ★ uniform_k_routing(·, 1) | +| `SAMPLE_ROUTING` | `(key, [B,T], step_f32) → tuple of routing draws`, statically sized; draws may be jointly sampled (independent repeats, antithetic/complementary subsets, per-step covers); `uniform_k` · `fixed_k` · `static_probability` · `scheduled_probability(p(step))` · `all` (S11) | ★ uniform_k_routing(·, 1) | | `SRC_STEP` | `adam(β₁,β₂,ε)` with bias correction; `sign` (`sources += lr·sign(grad)`) | ★ adam(.5, .99, 1e-8) | | `PROJ` / `EFFECTIVE` | **clamp**: PROJ = clamp[0,1], EFFECTIVE = identity, init U[0,1] · **sigmoid**: PROJ = identity (unbounded raw), EFFECTIVE = sigmoid, init N(0,1) | ★ clamp | | `SCOPE` | persistent: `c (1,1)` · `sc (1,T)` · `nsc (n,T), n|B` · `bsc (B,T)` (jax: `sc` + `bsc` today) — fresh-PGD: `c` · `bc` · `bsc` | ★ sc | diff --git a/param_decomp/configs.py b/param_decomp/configs.py index ad51a2958..a03f21329 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -46,6 +46,32 @@ class StaticProbabilityRoutingConfig(BaseConfig): p: Probability +class FixedKSubsetRoutingConfig(BaseConfig): + """Route each position to a uniformly-drawn subset of exactly `k` of the live sites.""" + + type: Literal["fixed_k_subset"] = "fixed_k_subset" + k: PositiveInt + + +class ScheduledProbabilityRoutingConfig(BaseConfig): + """`static_probability` with `p` evaluated per step from a schedule. The schedule ends + at `start_val * final_val_frac`; an increasing ramp uses `final_val_frac > 1`. Both + endpoints must be valid probabilities.""" + + type: Literal["scheduled_probability"] = "scheduled_probability" + p: ScheduleConfig + + @model_validator(mode="after") + def validate_probability_endpoints(self) -> Self: + end_val = self.p.start_val * self.p.final_val_frac + if not (self.p.start_val <= 1.0 and end_val <= 1.0): + raise ValueError( + f"p schedule endpoints must be probabilities: " + f"start={self.p.start_val}, end={end_val}" + ) + return self + + class AllRoutingConfig(BaseConfig): """Route every position to every module (the `"all"` fast path).""" @@ -53,7 +79,13 @@ class AllRoutingConfig(BaseConfig): # Discriminated union over the subset-routing configs (keyed by ``type``). -SubsetRoutingType = UniformKSubsetRoutingConfig | StaticProbabilityRoutingConfig | AllRoutingConfig +SubsetRoutingType = ( + UniformKSubsetRoutingConfig + | StaticProbabilityRoutingConfig + | FixedKSubsetRoutingConfig + | ScheduledProbabilityRoutingConfig + | AllRoutingConfig +) # --------------------------------------------------------------------------- @@ -428,7 +460,13 @@ class MergedStochasticPGDReconLossConfig(LossMetricConfig): canonical 0.5/0.5 pair). Carries the persistent-adversary fields; one source bundle feeds this one term (SPEC S23) and the S14' final ascent flows through its backward. A mixed assignment is a legal point of the mask box, so this samples the SAME - feasible set as the two-term objective, under a different (joint) sampler.""" + feasible set as the two-term objective, under a different (joint) sampler. + + `assignment` picks the Bernoulli granularity: `per_position` draws one bit per + (batch element, position) — a sequence's masked forward sees both families, coupled + through attention; `per_sample` draws one bit per batch element, so every position + in a sequence takes the same family (no within-sample mixing — each sample's loss is + scored against a pure-family context).""" type: Literal["MergedStochasticPGDReconLoss"] = "MergedStochasticPGDReconLoss" optimizer: AdamPGDConfig @@ -436,6 +474,7 @@ class MergedStochasticPGDReconLossConfig(LossMetricConfig): source_dtype: Literal["float32", "bfloat16"] = "float32" n_warmup_steps: NonNegativeInt = 0 adv_fraction: float = Field(gt=0.0, lt=1.0) + assignment: Literal["per_position", "per_sample"] = "per_position" routing: SubsetRoutingType = Field(default_factory=UniformKSubsetRoutingConfig) diff --git a/param_decomp/configs/pile4l_chunks_ab_1chunk_fk12_400k_cweast.yaml b/param_decomp/configs/pile4l_chunks_ab_1chunk_fk12_400k_cweast.yaml new file mode 100644 index 000000000..cac3e4a40 --- /dev/null +++ b/param_decomp/configs/pile4l_chunks_ab_1chunk_fk12_400k_cweast.yaml @@ -0,0 +1,182 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# CHUNK SWEEP rescue arm: single all-sites pass, fixed_k routing k=12 — an absolute +# routed-site budget per forward (deep-linear-identity finding: locality is a count, +# not a fraction; uniform-k routes ~n/2 sites). +run_name: pile4l-chunks-ab-1chunk-fk12-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: fixed_k_subset + k: 12 + type: StochasticReconSubsetLoss + - coeff: 0.5 + 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: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +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 diff --git a/param_decomp/configs/pile4l_chunks_ab_1chunk_fk3_400k_cweast.yaml b/param_decomp/configs/pile4l_chunks_ab_1chunk_fk3_400k_cweast.yaml new file mode 100644 index 000000000..86838d5be --- /dev/null +++ b/param_decomp/configs/pile4l_chunks_ab_1chunk_fk3_400k_cweast.yaml @@ -0,0 +1,182 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# CHUNK SWEEP rescue arm: single all-sites pass, fixed_k routing k=3 — an absolute +# routed-site budget per forward (deep-linear-identity finding: locality is a count, +# not a fraction; uniform-k routes ~n/2 sites). +run_name: pile4l-chunks-ab-1chunk-fk3-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: fixed_k_subset + k: 3 + type: StochasticReconSubsetLoss + - coeff: 0.5 + 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: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +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 diff --git a/param_decomp/configs/pile4l_chunks_ab_1chunk_fk6_400k_cweast.yaml b/param_decomp/configs/pile4l_chunks_ab_1chunk_fk6_400k_cweast.yaml new file mode 100644 index 000000000..46ba1ec19 --- /dev/null +++ b/param_decomp/configs/pile4l_chunks_ab_1chunk_fk6_400k_cweast.yaml @@ -0,0 +1,182 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# CHUNK SWEEP rescue arm: single all-sites pass, fixed_k routing k=6 — an absolute +# routed-site budget per forward (deep-linear-identity finding: locality is a count, +# not a fraction; uniform-k routes ~n/2 sites). +run_name: pile4l-chunks-ab-1chunk-fk6-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: fixed_k_subset + k: 6 + type: StochasticReconSubsetLoss + - coeff: 0.5 + 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: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +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 diff --git a/param_decomp/configs/pile4l_chunks_ab_1chunk_ramp_400k_cweast.yaml b/param_decomp/configs/pile4l_chunks_ab_1chunk_ramp_400k_cweast.yaml new file mode 100644 index 000000000..e8e4ae00e --- /dev/null +++ b/param_decomp/configs/pile4l_chunks_ab_1chunk_ramp_400k_cweast.yaml @@ -0,0 +1,184 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# CHUNK SWEEP rescue arm: single all-sites pass, scheduled-probability routing ramp +# 0.1 -> 1.0 linear (the routep-sched-1 champion on this harness under Lp). +run_name: pile4l-chunks-ab-1chunk-ramp-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: scheduled_probability + p: + start_val: 0.1 + fn_type: linear + final_val_frac: 10.0 + type: StochasticReconSubsetLoss + - coeff: 0.5 + 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: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +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 diff --git a/param_decomp/configs/pile4l_chunks_ab_chunks4_400k_cweast.yaml b/param_decomp/configs/pile4l_chunks_ab_chunks4_400k_cweast.yaml new file mode 100644 index 000000000..3fa450ca6 --- /dev/null +++ b/param_decomp/configs/pile4l_chunks_ab_chunks4_400k_cweast.yaml @@ -0,0 +1,180 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# CHUNK SWEEP interior point: 4 chunks (sites_per_chunk 6 = one layer per chunk). +run_name: pile4l-chunks-ab-chunks4-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + sites_per_chunk: 6 + type: ChunkwiseSubsetReconLoss + - coeff: 0.5 + 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: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +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 diff --git a/param_decomp/configs/pile4l_chunks_ab_chunks8_400k_cweast.yaml b/param_decomp/configs/pile4l_chunks_ab_chunks8_400k_cweast.yaml new file mode 100644 index 000000000..a4ab6f697 --- /dev/null +++ b/param_decomp/configs/pile4l_chunks_ab_chunks8_400k_cweast.yaml @@ -0,0 +1,182 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# CHUNK SWEEP reference: 8-chunk stochastic plan (sites_per_chunk 3 over the 24 sites, +# uniform-k within each chunk) — the 32L production plan shape staged at 4L. Everything +# else = the ab2 control harness. Task: recon-pass-consolidation-results. +run_name: pile4l-chunks-ab-chunks8-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + sites_per_chunk: 3 + type: ChunkwiseSubsetReconLoss + - coeff: 0.5 + 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: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +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 diff --git a/param_decomp/configs/pile4l_merge_ab2_control_400k_cweast.yaml b/param_decomp/configs/pile4l_merge_ab2_control_400k_cweast.yaml new file mode 100644 index 000000000..2f316ca88 --- /dev/null +++ b/param_decomp/configs/pile4l_merge_ab2_control_400k_cweast.yaml @@ -0,0 +1,179 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# ARM A (control): the stoch 0.5 + PPGD 0.5 pair, harness otherwise untouched. +run_name: pile4l-merge-ab2-control-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_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: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 16 +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 diff --git a/param_decomp/configs/pile4l_merge_ab2_merged_400k_cweast.yaml b/param_decomp/configs/pile4l_merge_ab2_merged_400k_cweast.yaml new file mode 100644 index 000000000..a8616ae2b --- /dev/null +++ b/param_decomp/configs/pile4l_merge_ab2_merged_400k_cweast.yaml @@ -0,0 +1,178 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# ARM B (merged, per-position): MergedStochasticPGDReconLoss replaces the pair. +run_name: pile4l-merge-ab2-merged-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - type: MergedStochasticPGDReconLoss + coeff: 1.0 + adv_fraction: 0.5 + routing: + type: uniform_k_subset + 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: bsc + n_warmup_steps: 2 + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 16 +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 diff --git a/param_decomp/configs/pile4l_merge_ab2_persample_400k_cweast.yaml b/param_decomp/configs/pile4l_merge_ab2_persample_400k_cweast.yaml new file mode 100644 index 000000000..e0a1ff5d6 --- /dev/null +++ b/param_decomp/configs/pile4l_merge_ab2_persample_400k_cweast.yaml @@ -0,0 +1,179 @@ +# Merge A/B round 2 on the CANONICAL p-20f9fc15 harness (Lucius's catch on the round-1 +# arms: their d1024x4 CI fn is ~8x smaller than Jose's d2048x8, and they swapped Lp +# imp-min for SmoothL0 — every round-1 arm, control included, sits well below the paper +# baseline: klCI 0.39-0.45 vs 0.337, klStoch 0.26-0.28 vs 0.21, PGD ~0.9-1.3 vs 0.64). +# This round changes ONLY the recon-term structure between arms; CI fn, seed, data, 400k +# steps all match pile_ppgd_bsc.yaml (the Jose replication) EXCEPT imp-min stays SmoothL0 +# (per Braun: earlier experiments showed SmoothL0 <-> Lp parity, and SmoothL0 is the +# going-forward penalty) - so only the CI-fn capacity changes vs round 1. +# data_files points at the crew user's staging of the same HF dataset (the canonical +# /mnt/data copy is group-locked). Task: recon-pass-consolidation-results. +# ARM C (merged, per-sample): as ARM B plus assignment: per_sample. +run_name: pile4l-merge-ab2-persample-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. DOCUMENTED EDITS vs upstream: +# 1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged +# parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 — +# the JAX trainer reads pre-tokenized parquet shards only). +# 2. data.eval_split: val -> train (parquet loading exposes a single "train" +# split; the staged _val dir exists for ad-hoc use but the loaders don't +# address it). +# 3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves +# nothing; the JAX leg checkpoints for requeue-resume + offline eval). +# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is +# upstream-identical. +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 10000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + type: chunkwise_transformer + blocks_per_chunk: 1 + d_model: 2048 + n_blocks: 8 + n_heads: 16 + mlp_hidden: 8192 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - type: MergedStochasticPGDReconLoss + coeff: 1.0 + adv_fraction: 0.5 + assignment: per_sample + routing: + type: uniform_k_subset + 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: bsc + n_warmup_steps: 2 + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 16 +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 diff --git a/param_decomp/configs/pile4l_merge_ab_merged_persample_400k_cweast.yaml b/param_decomp/configs/pile4l_merge_ab_merged_persample_400k_cweast.yaml new file mode 100644 index 000000000..9e295a931 --- /dev/null +++ b/param_decomp/configs/pile4l_merge_ab_merged_persample_400k_cweast.yaml @@ -0,0 +1,156 @@ +# PPGD-merge follow-up, arm B-PS (MERGED, per-sample assignment): identical to the +# merged arm except assignment: per_sample — one Bernoulli per batch element, so a +# sequence is all-adversarial or all-stochastic — no mask-family mixing within a sample +# (Braun, task recon-pass-consolidation-results). Same seed/data/CI as the original +# A/B arms for a paired comparison. +run_name: pile4l-merge-ab-merged-persample-400k +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + # pd-user-readable staging of the same HF dataset (danbraunai/pile-uncopyrighted-tok-shuffled); + # the canonical /mnt/data/.../pile_neox_tok_512 copy is group-locked to goodfire, which the + # crew user is not in. + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + 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: 1024 + n_blocks: 4 + n_heads: 8 + mlp_hidden: 4096 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - type: MergedStochasticPGDReconLoss + assignment: per_sample + coeff: 1.0 + adv_fraction: 0.5 + routing: + type: uniform_k_subset + 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: bsc + n_warmup_steps: 2 + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +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 diff --git a/param_decomp/configs/pile4l_merge_ab_merged_persample_400k_s1_cweast.yaml b/param_decomp/configs/pile4l_merge_ab_merged_persample_400k_s1_cweast.yaml new file mode 100644 index 000000000..6bf4ab5b8 --- /dev/null +++ b/param_decomp/configs/pile4l_merge_ab_merged_persample_400k_s1_cweast.yaml @@ -0,0 +1,156 @@ +# PPGD-merge follow-up, arm B-PS (MERGED, per-sample assignment): identical to the +# merged arm except assignment: per_sample — one Bernoulli per batch element, so a +# sequence is all-adversarial or all-stochastic — no mask-family mixing within a sample +# (Braun, task recon-pass-consolidation-results). Same seed/data/CI as the original +# A/B arms for a paired comparison — except seed 1 (seed-noise guard). +run_name: pile4l-merge-ab-merged-persample-400k-s1 +cadence: + keep_last_n_checkpoints: 2 + save_every: 5000 + train_log_every: 200 +data: + buffer_size: 1000 + column_name: input_ids + # pd-user-readable staging of the same HF dataset (danbraunai/pile-uncopyrighted-tok-shuffled); + # the canonical /mnt/data/.../pile_neox_tok_512 copy is group-locked to goodfire, which the + # crew user is not in. + data_files: /mnt/home/pd-user/param-decomp-runs/datasets/pile_neox_tok_512/data/*.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: c + 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: 1024 + n_blocks: 4 + n_heads: 8 + mlp_hidden: 4096 + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 400 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - frequency: + coeff: 0.0001 + reference_token_count: 65536 + coeff: 0.0002 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - type: MergedStochasticPGDReconLoss + assignment: per_sample + coeff: 1.0 + adv_fraction: 0.5 + routing: + type: uniform_k_subset + 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: bsc + n_warmup_steps: 2 + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 1 + steps: 400000 +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 diff --git a/param_decomp/experiments/invariance_check.py b/param_decomp/experiments/invariance_check.py index e30bf712f..059366749 100644 --- a/param_decomp/experiments/invariance_check.py +++ b/param_decomp/experiments/invariance_check.py @@ -120,6 +120,7 @@ def _run(steps: int, sharded: bool) -> list[dict[str, float]]: ppgd_cfg, ), lm.site_names, + 100, ) # fmt: skip step = make_train_step( lm=lm, diff --git a/param_decomp/recon.py b/param_decomp/recon.py index 6ddaf7584..0757869ef 100644 --- a/param_decomp/recon.py +++ b/param_decomp/recon.py @@ -28,6 +28,7 @@ CIMaskedReconLossConfig, CIMaskedReconSubsetLossConfig, FaithfulnessLossConfig, + FixedKSubsetRoutingConfig, ImportanceMinimalityLossConfig, MaskScopeLiteral, MergedStochasticPGDReconLossConfig, @@ -36,6 +37,7 @@ PGDReconLayerwiseLossConfig, PGDReconLossConfig, PGDReconSubsetLossConfig, + ScheduledProbabilityRoutingConfig, SmoothL0ImportanceMinimalityLossConfig, StaticProbabilityRoutingConfig, StochasticReconLayerwiseLossConfig, @@ -46,17 +48,20 @@ UnmaskedReconLossConfig, ) from param_decomp.lm import chunk_sites +from param_decomp.losses import scheduled_value_traced +from param_decomp.schedule import ScheduleConfig Routes = dict[str, Array] | None -RoutingSampler = Callable[[PRNGKeyArray, tuple[int, ...]], tuple[Routes, ...]] -"""`(key, leading_shape) -> (routes, ...)` — a STATICALLY-sized family of routing draws, -each `{site: bool[*leading]}` (or None = route everywhere) becoming ONE forward. The torch -`Router.get_masks` made pure: fresh draws per step require the key threaded in — -samplers run INSIDE the jitted step, so they must be traceable (SPEC R1). Returning +RoutingSampler = Callable[[PRNGKeyArray, tuple[int, ...], Array], tuple[Routes, ...]] +"""`(key, leading_shape, step_f32) -> (routes, ...)` — a STATICALLY-sized family of routing +draws, each `{site: bool[*leading]}` (or None = route everywhere) becoming ONE forward. The +torch `Router.get_masks` made pure: fresh draws per step require the key threaded in — +samplers run INSIDE the jitted step, so they must be traceable (SPEC R1). `step_f32` is the +traced step, read only by scheduled samplers (all others ignore it). Returning several draws from one invocation enables JOINTLY-sampled families (independent repeats, antithetic/complementary subsets, per-step random covers) that duplicated plan entries with independent keys cannot express. The plan's structure — live-sets, -sampler identities, family sizes — is static; only the key varies per step.""" +sampler identities, family sizes — is static; only the key (and step) vary per step.""" # ───────────────────────────── mask-source strategies ───────────────────────────── @@ -102,9 +107,10 @@ class PersistentSources: @dataclass(frozen=True) class MixedPersistentStochasticSources: - """The merged stochastic+PPGD strategy: per position, the persistent bundle's sources - (probability `cfg.adv_fraction`, routed all-live) or fresh `U[0,1]` (routed per the - entry's sampler). `state_key` indexes `TrainState.adversaries` like `PersistentSources`.""" + """The merged stochastic+PPGD strategy: per position (or per batch element, per + `cfg.assignment`), the persistent bundle's sources (probability `cfg.adv_fraction`, + routed all-live) or fresh `U[0,1]` (routed per the entry's sampler). `state_key` + indexes `TrainState.adversaries` like `PersistentSources`.""" state_key: str cfg: "MergedStochasticPGDReconLossConfig" @@ -211,7 +217,9 @@ def uniform_k_subset_routes( def uniform_k_routing(live_sites: tuple[str, ...], n_draws: int) -> RoutingSampler: """`n_draws` independent per-position uniform-k-subset draws over `live_sites`.""" - def sample(key: PRNGKeyArray, leading_shape: tuple[int, ...]) -> tuple[Routes, ...]: + def sample( + key: PRNGKeyArray, leading_shape: tuple[int, ...], _step_f32: Array + ) -> tuple[Routes, ...]: return tuple( uniform_k_subset_routes(draw_key, live_sites, leading_shape) for draw_key in random.split(key, n_draws) @@ -220,13 +228,60 @@ def sample(key: PRNGKeyArray, leading_shape: tuple[int, ...]) -> tuple[Routes, . return sample +def fixed_k_subset_routes( + key: PRNGKeyArray, live_sites: tuple[str, ...], k: int, leading_shape: tuple[int, ...] +) -> dict[str, Array]: + """Per position: a uniform subset of exactly `k` of the live sites routes True.""" + n_sites = len(live_sites) + perms = random.uniform(key, (n_sites, *leading_shape)).argsort(axis=0) + routed = perms < k + return {name: routed[j] for j, name in enumerate(live_sites)} + + +def fixed_k_routing(live_sites: tuple[str, ...], k: int, n_draws: int) -> RoutingSampler: + """`n_draws` independent per-position exactly-k-subset draws over `live_sites`.""" + assert 1 <= k <= len(live_sites), f"k={k} outside 1..{len(live_sites)} live sites" + + def sample( + key: PRNGKeyArray, leading_shape: tuple[int, ...], _step_f32: Array + ) -> tuple[Routes, ...]: + return tuple( + fixed_k_subset_routes(draw_key, live_sites, k, leading_shape) + for draw_key in random.split(key, n_draws) + ) + + return sample + + def static_probability_routing( live_sites: tuple[str, ...], p: float, n_draws: int ) -> RoutingSampler: """`n_draws` independent draws routing each position to each live site with probability `p` (torch `StaticProbabilityRouter`).""" - def sample(key: PRNGKeyArray, leading_shape: tuple[int, ...]) -> tuple[Routes, ...]: + def sample( + key: PRNGKeyArray, leading_shape: tuple[int, ...], _step_f32: Array + ) -> tuple[Routes, ...]: + return tuple( + { + name: random.bernoulli(random.fold_in(draw_key, j), p, leading_shape) + for j, name in enumerate(live_sites) + } + for draw_key in random.split(key, n_draws) + ) + + return sample + + +def scheduled_probability_routing( + live_sites: tuple[str, ...], p_schedule: ScheduleConfig, total_steps: int, n_draws: int +) -> RoutingSampler: + """`static_probability_routing` with `p` evaluated at the traced step.""" + + def sample( + key: PRNGKeyArray, leading_shape: tuple[int, ...], step_f32: Array + ) -> tuple[Routes, ...]: + p = scheduled_value_traced(step_f32, total_steps, p_schedule) return tuple( { name: random.bernoulli(random.fold_in(draw_key, j), p, leading_shape) @@ -241,20 +296,26 @@ def sample(key: PRNGKeyArray, leading_shape: tuple[int, ...]) -> tuple[Routes, . def route_all_n(n_draws: int) -> RoutingSampler: """`n_draws` forwards, each routing every position to every live site (`AllRoutingConfig`).""" - def sample(_key: PRNGKeyArray, _leading_shape: tuple[int, ...]) -> tuple[Routes, ...]: + def sample( + _key: PRNGKeyArray, _leading_shape: tuple[int, ...], _step_f32: Array + ) -> tuple[Routes, ...]: return (None,) * n_draws return sample def routing_sampler_from_config( - routing: SubsetRoutingType, live_sites: tuple[str, ...], n_draws: int + routing: SubsetRoutingType, live_sites: tuple[str, ...], n_draws: int, total_steps: int ) -> RoutingSampler: match routing: case UniformKSubsetRoutingConfig(): return uniform_k_routing(live_sites, n_draws) + case FixedKSubsetRoutingConfig(): + return fixed_k_routing(live_sites, routing.k, n_draws) case StaticProbabilityRoutingConfig(): return static_probability_routing(live_sites, routing.p, n_draws) + case ScheduledProbabilityRoutingConfig(): + return scheduled_probability_routing(live_sites, routing.p, total_steps, n_draws) case AllRoutingConfig(): return route_all_n(n_draws) @@ -289,15 +350,17 @@ def make_plan( routing: SubsetRoutingType, sources: MaskSourceStrategy, n_samples: int, + total_steps: int, ) -> ReconPlan: """One `ReconForward` per chunk: that chunk live, the rest frozen `x@W` (SPEC S2), with `n_samples` routing draws from `routing` over the chunk's own sites (SPEC S11) - and the shared `sources`. The chunking (`one_chunk`/`per_site`/`into_groups`) and the + and the shared `sources`. `total_steps` anchors scheduled routing (read only by + `scheduled_probability`). The chunking (`one_chunk`/`per_site`/`into_groups`) and the routing/source choices are orthogonal — see LOSS_PARITY_DESIGN.md.""" return tuple( ReconForward( live_sites=chunk, - sample_routing=routing_sampler_from_config(routing, chunk, n_samples), + sample_routing=routing_sampler_from_config(routing, chunk, n_samples, total_steps), sources=sources, ) for chunk in chunks @@ -309,11 +372,16 @@ def subset_chunk_plan( sites_per_chunk: int, n_samples: int, sources: MaskSourceStrategy, + total_steps: int, ) -> ReconPlan: """The production plan: `n_samples` uniform-k forwards per chunk (torch `SubsetReconPlan` over `ThreePoolTopology` chunks).""" return make_plan( - into_groups(site_names, sites_per_chunk), UniformKSubsetRoutingConfig(), sources, n_samples + into_groups(site_names, sites_per_chunk), + UniformKSubsetRoutingConfig(), + sources, + n_samples, + total_steps, ) @@ -339,6 +407,7 @@ def persistent_configs( def build_loss_terms( loss_metrics: Sequence[AnyLossMetricConfig], site_names: tuple[str, ...], + total_steps: int, ) -> LossSurface: """Validate the shared torch loss configs into the `LossSurface` record (LOSS_PARITY_DESIGN §3): exactly one faithfulness + one importance-minimality term, @@ -385,17 +454,21 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: case UnmaskedReconLossConfig() | CIMaskedReconLossConfig(): value = 1.0 if isinstance(cfg, UnmaskedReconLossConfig) else 0.0 plan = make_plan( - one_chunk(site_names), AllRoutingConfig(), ConstantSources(value), n_samples=1 + one_chunk(site_names), + AllRoutingConfig(), + ConstantSources(value), + 1, + total_steps, ) recon_terms.append(recon(cfg, plan)) case CIMaskedReconSubsetLossConfig(): plan = make_plan( - one_chunk(site_names), cfg.routing, ConstantSources(0.0), n_samples=1 + one_chunk(site_names), cfg.routing, ConstantSources(0.0), 1, total_steps ) recon_terms.append(recon(cfg, plan)) case CIMaskedReconLayerwiseLossConfig(): plan = make_plan( - per_site(site_names), AllRoutingConfig(), ConstantSources(0.0), n_samples=1 + per_site(site_names), AllRoutingConfig(), ConstantSources(0.0), 1, total_steps ) recon_terms.append(recon(cfg, plan)) case StochasticReconLossConfig(): @@ -404,11 +477,16 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: AllRoutingConfig(), StochasticSources(), cfg.n_mask_samples, + total_steps, ) recon_terms.append(recon(cfg, plan)) case StochasticReconSubsetLossConfig(): plan = make_plan( - one_chunk(site_names), cfg.routing, StochasticSources(), cfg.n_mask_samples + one_chunk(site_names), + cfg.routing, + StochasticSources(), + cfg.n_mask_samples, + total_steps, ) recon_terms.append(recon(cfg, plan)) case StochasticReconLayerwiseLossConfig(): @@ -417,6 +495,7 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: AllRoutingConfig(), StochasticSources(), cfg.n_mask_samples, + total_steps, ) recon_terms.append(recon(cfg, plan)) case ChunkwiseSubsetReconLossConfig(): @@ -426,6 +505,7 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: cfg.routing, StochasticSources(), cfg.n_samples, + total_steps, ) recon_terms.append(recon(cfg, plan)) case PGDReconLossConfig() | PGDReconSubsetLossConfig(): @@ -433,11 +513,11 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: routing = ( cfg.routing if isinstance(cfg, PGDReconSubsetLossConfig) else AllRoutingConfig() ) - plan = make_plan(one_chunk(site_names), routing, fresh, n_samples=1) + plan = make_plan(one_chunk(site_names), routing, fresh, 1, total_steps) recon_terms.append(recon(cfg, plan)) case PGDReconLayerwiseLossConfig(): fresh = FreshPGDSources(cfg.init, cfg.n_steps, cfg.step_size, cfg.mask_scope) - plan = make_plan(per_site(site_names), AllRoutingConfig(), fresh, n_samples=1) + plan = make_plan(per_site(site_names), AllRoutingConfig(), fresh, 1, total_steps) recon_terms.append(recon(cfg, plan)) case MergedStochasticPGDReconLossConfig(): schedule = cfg.optimizer.lr_schedule @@ -448,6 +528,7 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: cfg.routing, MixedPersistentStochasticSources(state_key=key, cfg=cfg), n_samples=1, + total_steps=total_steps, ) recon_terms.append(recon(cfg, plan)) case PersistentPGDReconLossConfig(): @@ -458,7 +539,8 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: one_chunk(site_names), AllRoutingConfig(), PersistentSources(state_key=key, cfg=cfg), - n_samples=1, + 1, + total_steps, ) recon_terms.append(recon(cfg, plan)) case _: diff --git a/param_decomp/run.py b/param_decomp/run.py index dae5d0e44..3474d3379 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -499,7 +499,7 @@ def run_decomposition_training( step_fn = make_train_step( lm=lm, - losses=build_loss_terms(pd.loss_metrics, lm.site_names), + losses=build_loss_terms(pd.loss_metrics, lm.site_names, pd.steps), components_optimizer=opt_vu, ci_fn_optimizer=opt_ci, total_steps=pd.steps, diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index ea7d019b8..c5751ec8b 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -118,7 +118,7 @@ def init_train_state( assert ci_fn.expects_axes == lm.leading_axes, ( f"CI fn expects leading axes {ci_fn.expects_axes} but model has {lm.leading_axes}" ) - losses = build_loss_terms(pd.loss_metrics, lm.site_names) + losses = build_loss_terms(pd.loss_metrics, lm.site_names, pd.steps) persistent = persistent_configs(losses.recon) term_coeff_by_state_key = { entry.sources.state_key: term.coeff diff --git a/param_decomp/tests/stacked_parity/test_stacked_parity.py b/param_decomp/tests/stacked_parity/test_stacked_parity.py index 4d05ec4d6..dc57caa18 100644 --- a/param_decomp/tests/stacked_parity/test_stacked_parity.py +++ b/param_decomp/tests/stacked_parity/test_stacked_parity.py @@ -195,7 +195,7 @@ def test_chunk_plan_static_live_set_matches(): f, lm, vu, resid = _load() plan = subset_chunk_plan( - lm.site_names, sites_per_chunk=3, n_samples=1, sources=StochasticSources() + lm.site_names, sites_per_chunk=3, n_samples=1, sources=StochasticSources(), total_steps=100 ) chunk0 = lm.site_names[:3] assert plan[0].live_sites == chunk0, (plan[0].live_sites, chunk0) @@ -277,6 +277,7 @@ def test_train_trajectory_matches(): ppgd_cfg, ), lm.site_names, + 100, ) step_fn = make_train_step( lm=lm, diff --git a/param_decomp/tests/test_checkpoint.py b/param_decomp/tests/test_checkpoint.py index ea4cddbf2..168772f87 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -128,6 +128,7 @@ def _build(seed: int): ppgd_cfg, ), lm.site_names, + 100, ) # fmt: skip step = make_train_step( lm=lm, diff --git a/param_decomp/tests/test_checkpoint_production_topology.py b/param_decomp/tests/test_checkpoint_production_topology.py index 85b3ef3a3..5f8837964 100644 --- a/param_decomp/tests/test_checkpoint_production_topology.py +++ b/param_decomp/tests/test_checkpoint_production_topology.py @@ -155,6 +155,7 @@ def _build_sharded(seed: int): *ppgd_cfgs, ), lm.site_names, + 100, ) # fmt: skip assert tuple(persistent_configs(loss_terms.recon)) == PERSISTENT_TERMS, loss_terms diff --git a/param_decomp/tests/test_config.py b/param_decomp/tests/test_config.py index 534762d73..8a14261b7 100644 --- a/param_decomp/tests/test_config.py +++ b/param_decomp/tests/test_config.py @@ -96,7 +96,9 @@ def test_b128_config_converts(): assert converted.data is not None and converted.data.global_batch == 128 assert converted.target.sites == mlp_family_site_cs(18, 18, 24576) losses = build_loss_terms( - converted.pd.loss_metrics, tuple(sc.name for sc in converted.target.sites) + converted.pd.loss_metrics, + tuple(sc.name for sc in converted.target.sites), + converted.pd.steps, ) faith, imp = losses.faith, losses.imp assert isinstance(imp.cfg, ImportanceMinimalityLossConfig) diff --git a/param_decomp/tests/test_generic_model_io.py b/param_decomp/tests/test_generic_model_io.py index 7640df7b3..c1919f056 100644 --- a/param_decomp/tests/test_generic_model_io.py +++ b/param_decomp/tests/test_generic_model_io.py @@ -259,6 +259,7 @@ def test_train_step_runs_through_generic_target(): StochasticReconLossConfig(coeff=1.0), ), lm.site_names, + 10, ) step_fn = make_train_step( lm=lm, diff --git a/param_decomp/tests/test_llama8b.py b/param_decomp/tests/test_llama8b.py index a22f6ed80..7092f1aa6 100644 --- a/param_decomp/tests/test_llama8b.py +++ b/param_decomp/tests/test_llama8b.py @@ -385,6 +385,7 @@ def test_step_trains_and_has_vpd_signature(site_cs: tuple[SiteC, ...]): ppgd_cfg, ), lm.site_names, + 100, ) step = make_train_step( lm=lm, @@ -506,6 +507,7 @@ def run_step(n_ascent_steps: int) -> tuple[TrainState, dict[str, jax.Array]]: ), ), lm.site_names, + 100, ) step = make_train_step( lm=lm, diff --git a/param_decomp/tests/test_llama_simple_mlp.py b/param_decomp/tests/test_llama_simple_mlp.py index 567e524ea..b7af742e8 100644 --- a/param_decomp/tests/test_llama_simple_mlp.py +++ b/param_decomp/tests/test_llama_simple_mlp.py @@ -403,6 +403,7 @@ def test_step_trains_and_has_vpd_signature(): ppgd_cfg, ), lm.site_names, + 100, ) step = make_train_step( lm=lm, @@ -475,6 +476,18 @@ def test_decomp_vu_shapes_fp32(): _REAL_CACHE_DIR = Path("/mnt/data/artifacts/mechanisms/param-decomp/pretrain_cache/spd-t-9d2b8f02") + + +def _real_cache_mounted() -> bool: + # Path.exists() raises PermissionError (not False) when a parent dir denies stat — + # the cache lives under a group-restricted dir, so collection dies for users + # outside that group without this guard. + try: + return _REAL_CACHE_DIR.exists() + except PermissionError: + return False + + _PRODUCTION_PATTERN_CS = { "h.*.mlp.c_fc": 3072, "h.*.mlp.down_proj": 3584, @@ -486,7 +499,7 @@ def test_decomp_vu_shapes_fp32(): """The pile production decomposition (torch `pile_llama_simple_mlp-4L.yaml`).""" -@pytest.mark.skipif(not _REAL_CACHE_DIR.exists(), reason="t-9d2b8f02 pretrain cache not mounted") +@pytest.mark.skipif(not _real_cache_mounted(), reason="t-9d2b8f02 pretrain cache not mounted") def test_pretrained_target_converts_with_wildcards(): """`kind: pretrained` LlamaSimpleMLP target specs convert, expanding `h.*` wildcard decomposition patterns over the checkpoint's n_layer (4).""" @@ -530,6 +543,7 @@ def test_pretrained_target_converts_with_wildcards(): loss_terms = build_loss_terms( cfg.pd.loss_metrics, tuple(sc.name for sc in target.sites), + cfg.pd.steps, ) (stoch_term,) = [t for t in loss_terms.recon if t.name == "StochasticReconSubsetLoss"] (stoch_entry,) = stoch_term.plan diff --git a/param_decomp/tests/test_merged_recon.py b/param_decomp/tests/test_merged_recon.py index 5012a3dcc..6430db3fa 100644 --- a/param_decomp/tests/test_merged_recon.py +++ b/param_decomp/tests/test_merged_recon.py @@ -1,9 +1,12 @@ """MergedStochasticPGDReconLoss: the one-forward stoch+PPGD term through the jitted step.""" +from typing import Literal + import equinox as eqx import jax import jax.numpy as jnp import optax +import pytest from param_decomp.adversary import ( PersistentAdversary, @@ -31,10 +34,13 @@ from param_decomp.train import TrainState, make_train_step -def _merged_cfg(n_warmup: int) -> MergedStochasticPGDReconLossConfig: +def _merged_cfg( + n_warmup: int, assignment: Literal["per_position", "per_sample"] = "per_position" +) -> MergedStochasticPGDReconLossConfig: return MergedStochasticPGDReconLossConfig( coeff=1.0, adv_fraction=0.5, + assignment=assignment, routing=UniformKSubsetRoutingConfig(), scope=SCScope(), optimizer=AdamPGDConfig( @@ -56,6 +62,7 @@ def test_merged_term_builds_one_entry(): cfg, ), ("a", "b"), + total_steps=100, ) (term,) = losses.recon (entry,) = term.plan @@ -64,7 +71,8 @@ def test_merged_term_builds_one_entry(): assert entry.has_delta -def test_merged_train_step_end_to_end(): +@pytest.mark.parametrize("assignment", ["per_position", "per_sample"]) +def test_merged_train_step_end_to_end(assignment: Literal["per_position", "per_sample"]): """Full jitted step with ONE merged recon term: finite losses, the persistent adversary updates (n_warmup + 1 per step through warmup + the S14' final ascent), sources stay projected.""" @@ -78,7 +86,7 @@ def test_merged_train_step_end_to_end(): opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0)) opt_ci = optax.adamw(1e-3, weight_decay=0.0) - merged = _merged_cfg(n_warmup) + merged = _merged_cfg(n_warmup, assignment) src = init_persistent_sources( lm.site_names, tuple(s.C for s in lm.sites), (1, seq), jnp.float32, jax.random.PRNGKey(3) ) @@ -109,6 +117,7 @@ def test_merged_train_step_end_to_end(): merged, ), lm.site_names, + total_steps=100, ) step = make_train_step( lm=lm, diff --git a/param_decomp/tests/test_no_bake_invariant.py b/param_decomp/tests/test_no_bake_invariant.py index 92f809081..a594d21c2 100644 --- a/param_decomp/tests/test_no_bake_invariant.py +++ b/param_decomp/tests/test_no_bake_invariant.py @@ -80,6 +80,7 @@ def _build_step_and_args(): StochasticReconLossConfig(coeff=1.0), ), lm.site_names, + 10, ) step_fn = make_train_step( lm=lm, diff --git a/param_decomp/tests/test_recon_log_keys.py b/param_decomp/tests/test_recon_log_keys.py index 5025b976f..4a58f89b6 100644 --- a/param_decomp/tests/test_recon_log_keys.py +++ b/param_decomp/tests/test_recon_log_keys.py @@ -76,6 +76,7 @@ def _recon_terms(recon_configs: tuple[object, ...]) -> tuple[ReconLossTerm, ...] terms = build_loss_terms( (*_non_recon_configs(), *recon_configs), # pyright: ignore[reportArgumentType] site_names=SITE_NAMES, + total_steps=100, ) return terms.recon diff --git a/param_decomp/tests/test_routing.py b/param_decomp/tests/test_routing.py new file mode 100644 index 000000000..2480e415b --- /dev/null +++ b/param_decomp/tests/test_routing.py @@ -0,0 +1,83 @@ +"""Routing samplers: per-position subset laws and the scheduled-p step dependence (S11).""" + +import jax.numpy as jnp +import pytest +from jax import random + +from param_decomp.configs import ( + FixedKSubsetRoutingConfig, + ScheduledProbabilityRoutingConfig, + StaticProbabilityRoutingConfig, + UniformKSubsetRoutingConfig, +) +from param_decomp.recon import routing_sampler_from_config +from param_decomp.schedule import ScheduleConfig, get_scheduled_value + +SITES = tuple(f"s{i}" for i in range(8)) +LEADING = (16, 32) +STEP0 = jnp.zeros((), jnp.float32) + + +def _routed_count_per_position(routes: dict[str, jnp.ndarray]) -> jnp.ndarray: + return jnp.stack([routes[s] for s in SITES]).sum(axis=0) + + +@pytest.mark.parametrize("k", [1, 3, 8]) +def test_fixed_k_routes_exactly_k_per_position(k: int): + sampler = routing_sampler_from_config( + FixedKSubsetRoutingConfig(k=k), SITES, n_draws=2, total_steps=100 + ) + for routes in sampler(random.PRNGKey(0), LEADING, STEP0): + assert routes is not None + assert bool((_routed_count_per_position(routes) == k).all()) + + +def test_fixed_k_refuses_k_above_n_sites(): + with pytest.raises(AssertionError): + routing_sampler_from_config( + FixedKSubsetRoutingConfig(k=9), SITES, n_draws=1, total_steps=100 + ) + + +def test_uniform_k_spans_1_to_n_per_position(): + sampler = routing_sampler_from_config( + UniformKSubsetRoutingConfig(), SITES, n_draws=1, total_steps=100 + ) + (routes,) = sampler(random.PRNGKey(0), LEADING, STEP0) + assert routes is not None + counts = _routed_count_per_position(routes) + assert counts.min() >= 1 and counts.max() <= len(SITES) + + +def test_static_probability_matches_p(): + sampler = routing_sampler_from_config( + StaticProbabilityRoutingConfig(p=0.25), SITES, n_draws=1, total_steps=100 + ) + (routes,) = sampler(random.PRNGKey(0), (64, 64), STEP0) + assert routes is not None + mean = jnp.stack([routes[s] for s in SITES]).mean() + assert abs(float(mean) - 0.25) < 0.02 + + +def test_scheduled_probability_tracks_schedule(): + total_steps = 100 + schedule = ScheduleConfig(start_val=0.1, fn_type="linear", final_val_frac=10.0) + sampler = routing_sampler_from_config( + ScheduledProbabilityRoutingConfig(p=schedule), SITES, n_draws=1, total_steps=total_steps + ) + for step in (0, 50, 99): + expected = get_scheduled_value(step, total_steps, schedule) + (routes,) = sampler(random.PRNGKey(0), (64, 64), jnp.float32(step)) + assert routes is not None + mean = float(jnp.stack([routes[s] for s in SITES]).mean()) + assert abs(mean - expected) < 0.02, (step, mean, expected) + assert abs(get_scheduled_value(99, total_steps, schedule) - 1.0) < 1e-6 + + +def test_scheduled_probability_refuses_endpoints_above_1(): + with pytest.raises(ValueError): + ScheduledProbabilityRoutingConfig( + p=ScheduleConfig(start_val=0.5, fn_type="linear", final_val_frac=4.0) + ) + with pytest.raises(ValueError): + ScheduledProbabilityRoutingConfig(p=ScheduleConfig(start_val=1.5, fn_type="constant")) diff --git a/param_decomp/train.py b/param_decomp/train.py index 1dc940f46..3fddb83ec 100644 --- a/param_decomp/train.py +++ b/param_decomp/train.py @@ -321,7 +321,7 @@ def warmup_scoring_loss(sources: dict[str, Array]) -> Array: continue fresh_cfg = entry.sources routing_key, init_key = random.split(random.fold_in(term_key, entry_idx)) - routes_per_draw = entry.sample_routing(routing_key, leading) + routes_per_draw = entry.sample_routing(routing_key, leading, step_f32) fixed_routes[(term_idx, entry_idx)] = routes_per_draw live_specs = tuple(s for s in sites if s.name in entry.live_sites) init = init_fresh_pgd_sources( @@ -394,7 +394,7 @@ def loss_fn( case FreshPGDSources(): routes_per_draw = fixed_routes[(term_idx, entry_idx)] case _: - routes_per_draw = entry.sample_routing(routing_key, leading) + routes_per_draw = entry.sample_routing(routing_key, leading, step_f32) for draw_idx, routes in enumerate(routes_per_draw): draw_key = random.fold_in(entry_key, draw_idx) @@ -450,15 +450,26 @@ def pre_built_fwd( ) ) case MixedPersistentStochasticSources(state_key=state_key): - # Per-position mixing (SPEC S10' variation): sources are - # the persistent bundle's where `assign`, fresh U[0,1] + # Mixing (SPEC S10' variation): sources are the + # persistent bundle's where `assign`, fresh U[0,1] # elsewhere; adversarial positions route all-live, the # rest keep this draw's routing. One forward serves both # pressures; the persistent sources stay live leaves so # the S14' final ascent rides this term's backward. + # `assign` is per (batch, position) or — with + # assignment=per_sample — per batch element, shape + # (B, 1, ...) broadcasting over the position axes so a + # sequence never mixes families. mix_cfg = entry.sources.cfg akey, ukey = random.split(draw_key) - assign = random.bernoulli(akey, mix_cfg.adv_fraction, leading) + assign_shape = ( + leading + if mix_cfg.assignment == "per_position" + else (leading[0], *(1,) * (len(leading) - 1)) + ) + assign = random.bernoulli( + akey, mix_cfg.adv_fraction, assign_shape + ) fresh_uniform = { site: random.uniform( random.fold_in(ukey, site_idx), diff --git a/param_decomp_lab/experiments/lm/config.py b/param_decomp_lab/experiments/lm/config.py index b70313ce6..e98199787 100644 --- a/param_decomp_lab/experiments/lm/config.py +++ b/param_decomp_lab/experiments/lm/config.py @@ -306,7 +306,7 @@ def _assert_losses_supported(cfg: LMExperimentConfig, site_names: tuple[str, ... """Run the schema's loss configs through `build_loss_terms` so unsupported metrics refuse at convert time rather than on the GPUs. The engine reads `pd.loss_metrics` verbatim (yaml order is RNG-load-bearing), so nothing is returned.""" - build_loss_terms(cfg.pd.loss_metrics, site_names) + build_loss_terms(cfg.pd.loss_metrics, site_names, cfg.pd.steps) def _data(cfg: LMExperimentConfig) -> DataConfig: diff --git a/param_decomp_lab/experiments/resid_mlp/run.py b/param_decomp_lab/experiments/resid_mlp/run.py index 526296395..ff1ffb523 100644 --- a/param_decomp_lab/experiments/resid_mlp/run.py +++ b/param_decomp_lab/experiments/resid_mlp/run.py @@ -50,6 +50,7 @@ def build_resid_mlp_built_run(cfg: ResidMLPExperimentConfig, run_id: str) -> Bui build_loss_terms( cfg.pd.loss_metrics, tuple(sc.name for sc in site_cs), + cfg.pd.steps, ) target = resid_mlp.ResidMLPTargetConfig( n_features=cfg.target.n_features, diff --git a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py index 3545eff3c..cf3bdf173 100644 --- a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py +++ b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py @@ -252,7 +252,7 @@ def _make_state_and_step( ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)), adversaries={}, step=jnp.zeros((), jnp.int32), ) # fmt: skip - loss_terms = build_loss_terms(_loss_metrics(), lm.site_names) + loss_terms = build_loss_terms(_loss_metrics(), lm.site_names, 100) step = make_train_step( lm=lm, losses=loss_terms, components_optimizer=opt_vu, ci_fn_optimizer=opt_ci, total_steps=total_steps, remat_recon_forwards=False, remat_ci_fn=False, mesh=None, @@ -358,7 +358,7 @@ def _faith_warmed_state( ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)), adversaries={}, step=jnp.zeros((), jnp.int32), ) # fmt: skip - loss_terms = build_loss_terms(_recovery_loss_metrics(), lm.site_names) + loss_terms = build_loss_terms(_recovery_loss_metrics(), lm.site_names, 100) step = make_train_step( lm=lm, losses=loss_terms, components_optimizer=opt_vu, ci_fn_optimizer=opt_ci, total_steps=total_steps, remat_recon_forwards=False, remat_ci_fn=False, mesh=None, diff --git a/param_decomp_lab/experiments/tms/run.py b/param_decomp_lab/experiments/tms/run.py index 56e60a80e..cb8b6593f 100644 --- a/param_decomp_lab/experiments/tms/run.py +++ b/param_decomp_lab/experiments/tms/run.py @@ -52,6 +52,7 @@ def build_tms_built_run(cfg: TMSExperimentConfig, run_id: str) -> BuiltRun: build_loss_terms( cfg.pd.loss_metrics, tuple(sc.name for sc in site_cs), + cfg.pd.steps, ) target = tms.TMSTargetConfig( n_features=cfg.target.n_features, diff --git a/param_decomp_lab/experiments/tms/test_tms.py b/param_decomp_lab/experiments/tms/test_tms.py index 89932c322..d643aa3a5 100644 --- a/param_decomp_lab/experiments/tms/test_tms.py +++ b/param_decomp_lab/experiments/tms/test_tms.py @@ -202,7 +202,7 @@ def _make_state_and_step( ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)), adversaries={}, step=jnp.zeros((), jnp.int32), ) # fmt: skip - loss_terms = build_loss_terms(_loss_metrics(), lm.site_names) + loss_terms = build_loss_terms(_loss_metrics(), lm.site_names, 100) step = make_train_step( lm=lm, losses=loss_terms, components_optimizer=opt_vu, ci_fn_optimizer=opt_ci, total_steps=total_steps, remat_recon_forwards=False, remat_ci_fn=False, mesh=None, @@ -302,7 +302,7 @@ def _faith_warmed_state( ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)), adversaries={}, step=jnp.zeros((), jnp.int32), ) # fmt: skip - loss_terms = build_loss_terms(_recovery_loss_metrics(), lm.site_names) + loss_terms = build_loss_terms(_recovery_loss_metrics(), lm.site_names, 100) step = make_train_step( lm=lm, losses=loss_terms, components_optimizer=opt_vu, ci_fn_optimizer=opt_ci, total_steps=total_steps, remat_recon_forwards=False, remat_ci_fn=False, mesh=None,