diff --git a/CLAUDE.md b/CLAUDE.md index fae0420fb..1a42aaa86 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -174,9 +174,9 @@ and returns JAX-native as the #10 torch->jax adapter. (the LM composition root — `python -m param_decomp_lab.experiments.lm.run`), `config.py` (LM schema + LM build), `load_run.py` (open a finished JAX run), `data.py` / `prestage_tokenized.py` (offline tokenize → parquet shards), `jax_launch.py` (`pd-lm`). - The TMS and ResidualMLP domains live under `experiments/{tms,resid_mlp}/` - (`run.py` + `config.py` + `model.py`; `pd-tms` / `pd-resid-mlp`), calling the core - engine as a library. + The TMS, ResidualMLP and DeepLinear domains live under + `experiments/{tms,resid_mlp,deep_linear}/` (`run.py` + `config.py` + `model.py`; + `pd-tms` / `pd-resid-mlp` / `pd-deep-linear`), calling the core engine as a library. - `param_decomp_lab/{harvest,autointerp,clustering,investigate}/` — post-pipeline stages, each with its own CLAUDE.md. - `param_decomp_lab/postprocess/` — orchestrates the post-pipeline stages. @@ -252,7 +252,7 @@ via `pd-lm`. Slow/plot eval is in-loop only (no CLI). | `python -m pretrain.train` | `pretrain/train.py` | The core in-house target-LM pretrainer | | `pd-lm` | `experiments/lm/launch.py` | Launch a decomposition trainer run; config-driven via `runtime.dp` (`dp=N` → snapshot + pinned launch config + sbatch across `N//8` nodes, per-node job-side venv; `dp=null` → inline) | | `pd-pretrain` | `experiments/lm/pretrain/launch.py` | Launch a pretrainer run; config-driven via `dp` (`dp=N` → sbatch; `dp=null` → inline) | -| `pd-tms` / `pd-resid-mlp` | `experiments/{tms,resid_mlp}/run.py` | The CPU toy decomposition CLIs | +| `pd-tms` / `pd-resid-mlp` / `pd-deep-linear` | `experiments/{tms,resid_mlp,deep_linear}/run.py` | The CPU toy decomposition CLIs | | `pd-harvest` | `harvest/scripts/run_slurm_cli.py` | Submit harvest SLURM job | | `pd-autointerp` | `autointerp/scripts/run_slurm_cli.py` | Submit autointerp SLURM job | | `pd-clustering` / `pd-cluster-merge` / `pd-cluster-distances` | `clustering/scripts/` | Clustering ensemble / merge / consensus distances | diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index d1a30ab3e..a640e6f96 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 455375223..92632d5cc 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 +) # --------------------------------------------------------------------------- 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 6624788dd..195544a08 100644 --- a/param_decomp/recon.py +++ b/param_decomp/recon.py @@ -28,6 +28,7 @@ CIMaskedReconLossConfig, CIMaskedReconSubsetLossConfig, FaithfulnessLossConfig, + FixedKSubsetRoutingConfig, ImportanceMinimalityLossConfig, MaskScopeLiteral, PersistentPGDReconLossConfig, @@ -35,6 +36,7 @@ PGDReconLayerwiseLossConfig, PGDReconLossConfig, PGDReconSubsetLossConfig, + ScheduledProbabilityRoutingConfig, SmoothL0ImportanceMinimalityLossConfig, StaticProbabilityRoutingConfig, StochasticReconLayerwiseLossConfig, @@ -45,17 +47,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 ───────────────────────────── @@ -194,7 +199,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) @@ -203,13 +210,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) @@ -224,20 +278,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) @@ -272,15 +332,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 @@ -292,11 +354,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, ) @@ -321,6 +388,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, @@ -367,17 +435,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(): @@ -386,11 +458,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(): @@ -399,15 +476,16 @@ def recon(cfg: AnyLossMetricConfig, plan: ReconPlan) -> ReconLossTerm: AllRoutingConfig(), StochasticSources(), cfg.n_mask_samples, + total_steps, ) recon_terms.append(recon(cfg, plan)) case ChunkwiseSubsetReconLossConfig(): - assert isinstance(cfg.routing, UniformKSubsetRoutingConfig), cfg.routing plan = make_plan( into_groups(site_names, cfg.sites_per_chunk), cfg.routing, StochasticSources(), cfg.n_samples, + total_steps, ) recon_terms.append(recon(cfg, plan)) case PGDReconLossConfig() | PGDReconSubsetLossConfig(): @@ -415,11 +493,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 PersistentPGDReconLossConfig(): schedule = cfg.optimizer.lr_schedule @@ -429,7 +507,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 17ebca0eb..39e24e876 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -508,7 +508,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 a20bfc436..8b99084ce 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -117,7 +117,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 38a0519dd..cc7fee838 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 00fb594b2..6bd6e8e4c 100644 --- a/param_decomp/tests/test_llama8b.py +++ b/param_decomp/tests/test_llama8b.py @@ -412,6 +412,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, @@ -533,6 +534,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..e312b76d7 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, @@ -530,6 +531,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_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 6694d2157..af7ef94c7 100644 --- a/param_decomp/train.py +++ b/param_decomp/train.py @@ -319,7 +319,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( @@ -392,7 +392,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) diff --git a/param_decomp_lab/experiments/CLAUDE.md b/param_decomp_lab/experiments/CLAUDE.md index 246e41758..b2f3c57cb 100644 --- a/param_decomp_lab/experiments/CLAUDE.md +++ b/param_decomp_lab/experiments/CLAUDE.md @@ -14,7 +14,7 @@ autointerp/clustering read a run's target topology from `experiments.lm.load_run.run_metadata` (config + pretrain cache, no checkpoint restore) — see `param_decomp_lab/adapters/pd.py`. -## Toy domains (TMS, ResidMLP) +## Toy domains (TMS, ResidMLP, DeepLinear) The TMS and ResidualMLP toys are LAB experiments that call the core engine as a library (the core itself has zero toy-specific code). Each `experiments/{tms,resid_mlp}/` carries: @@ -38,6 +38,14 @@ The TMS and ResidualMLP toys are LAB experiments that call the core engine as a - `configs/*.yaml` — the canonical `experiments.{tms,resid_mlp}.config` schema (TMS: 5-2 / 40-10 / the `-id` deeper variants; ResidMLP: 1l/2l/3l + the global-CI variant). +`experiments/deep_linear/` is the third toy, same shape minus pretraining: `n_layers` +FROZEN `eye(n_features)` sites (`layers.{i}`, `layers.*` wildcard expansion), uniform +one-hot data sampled on the fly, and — unlike the MSE toys — the LM's `kl_per_position` +recon on the final logits, making it the smallest end-to-end `ChunkwiseSubsetReconLoss` +testbed (`sites_per_chunk` must divide `n_layers`). No `PersistentPGDReconLoss`: the +engine pins persistent sources to a sequence axis, and positionless targets have none +(`build_deep_linear_built_run` rejects it loudly). + TMS deeper variant (`n_hidden_layers>0`, the `-id` configs) + the ResidMLP `global` CI arch (`fn_type=global_shared_mlp`) are restored and wired end-to-end (the global arch dispatches through the core `init_train_state` via `experiments.config.ci_arch`). Toy harvest / @@ -61,7 +69,8 @@ experiments/ │ ├── prestage_tokenized.py # HF text -> int32 parquet shards for the JAX trainer │ └── arithmetic_probe.py # a x b arithmetic grid spec -> in-memory eval probe (ArithmeticCIGrid) ├── tms/ # pd-tms (CPU): model.py + run.py + configs/ + test_tms.py -└── resid_mlp/ # pd-resid-mlp (CPU): model.py + run.py + configs/ + test +├── resid_mlp/ # pd-resid-mlp (CPU): model.py + run.py + configs/ + test +└── deep_linear/ # pd-deep-linear (CPU/1 GPU): frozen identity stack, KL recon, chunkwise-recon testbed ``` ## LM `target.spec` diff --git a/param_decomp_lab/experiments/deep_linear/__init__.py b/param_decomp_lab/experiments/deep_linear/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/param_decomp_lab/experiments/deep_linear/config.py b/param_decomp_lab/experiments/deep_linear/config.py new file mode 100644 index 000000000..e6428aec5 --- /dev/null +++ b/param_decomp_lab/experiments/deep_linear/config.py @@ -0,0 +1,30 @@ +"""Deep-linear identity experiment config schema — the target is fully determined by +`(n_features, n_layers)` (identity weights, constructed not pretrained) and the data is +uniform one-hot rows (no knobs), so a run is reproducible from the config alone.""" + +from typing import Literal + +from pydantic import PositiveInt + +from param_decomp.base_config import BaseConfig +from param_decomp_lab.experiments.config import ExperimentConfig + + +class DeepLinearTargetConfig(BaseConfig): + """`n_layers` frozen `eye(n_features)` sites named `layers.{i}`. `logit_scale` + multiplies the final output only (softmax temperature for the KL recon; layers stay + exact identities); `recon` picks the comparison (`kl` = LM semantics, `mse` = raw + logits, no softmax saturation).""" + + n_features: PositiveInt + n_layers: PositiveInt + logit_scale: float = 1.0 + recon: Literal["kl", "mse"] = "kl" + + +class DeepLinearDataConfig(BaseConfig): + """Uniform one-hot rows; nothing to configure.""" + + +class DeepLinearExperimentConfig(ExperimentConfig[DeepLinearTargetConfig, DeepLinearDataConfig]): + pass diff --git a/param_decomp_lab/experiments/deep_linear/configs/deep_linear_128.yaml b/param_decomp_lab/experiments/deep_linear/configs/deep_linear_128.yaml new file mode 100644 index 000000000..189f55411 --- /dev/null +++ b/param_decomp_lab/experiments/deep_linear/configs/deep_linear_128.yaml @@ -0,0 +1,68 @@ +# Deep-linear identity decomposition, d=128 — the smallest target that exercises +# ChunkwiseSubsetReconLoss end-to-end. Every site is a frozen eye(128); with one-hot +# inputs the ground truth is 128 live rank-1 components e_i e_iT per site, so +# eval/identity_ci_error/ -> 0 is full recovery. Loss stack mirrors the production +# pile_ppgd_bsc.yaml minus the persistent-PGD adversary (the engine pins persistent +# sources to a sequence axis, which positionless targets lack), at toy-scale LRs. layers.* expands to +# n_layers sites; the sweep rewrites n_layers (5/50/200) and sites_per_chunk (1 = one +# chunk per layer ... n_layers = 1 chunk == StochasticReconSubsetLoss; must divide +# n_layers). Runs on 1 GPU or CPU. +run_name: jax-deep-linear-128-L5 +pd: + seed: 0 + ci_config: + type: layerwise_mlp + hidden_dims: + - 128 + decomposition_targets: + - module_pattern: layers.* + C: 256 + loss_metrics: + - type: ImportanceMinimalityLoss + coeff: 2e-4 + eps: 1e-6 + pnorm: + start_val: 2.0 + fn_type: linear + final_val_frac: 0.2 + frequency: + coeff: 1e-4 + reference_token_count: 1024 + - type: ChunkwiseSubsetReconLoss + coeff: 1.0 + sites_per_chunk: 1 + routing: + type: uniform_k_subset + - type: FaithfulnessLoss + coeff: 1.0 + batch_size: 1024 + steps: 20000 + components_optimizer: + grad_clip_norm: 0.01 + lr_schedule: + start_val: 1e-3 + fn_type: cosine + warmup_pct: 0.0 + final_val_frac: 0.1 + ci_fn_optimizer: + lr_schedule: + start_val: 1e-3 + fn_type: cosine + warmup_pct: 0.0 + final_val_frac: 0.1 + faithfulness_warmup_steps: 200 + faithfulness_warmup_lr: 0.01 + faithfulness_warmup_weight_decay: 0.0 +runtime: + remat_recon_forwards: false + dp: null +cadence: + train_log_every: 200 + save_every: 5000 + keep_last_n_checkpoints: 2 +target: + n_features: 128 + n_layers: 5 +data: {} +wandb: + project: param-decomp diff --git a/param_decomp_lab/experiments/deep_linear/configs/deep_linear_30.yaml b/param_decomp_lab/experiments/deep_linear/configs/deep_linear_30.yaml new file mode 100644 index 000000000..fa8fba12a --- /dev/null +++ b/param_decomp_lab/experiments/deep_linear/configs/deep_linear_30.yaml @@ -0,0 +1,68 @@ +# Deep-linear identity decomposition, d=30 — the smallest target that exercises +# ChunkwiseSubsetReconLoss end-to-end. Every site is a frozen eye(30); with one-hot +# inputs the ground truth is 30 live rank-1 components e_i e_iT per site, so +# eval/identity_ci_error/ -> 0 is full recovery. Loss stack mirrors the production +# pile_ppgd_bsc.yaml minus the persistent-PGD adversary (the engine pins persistent +# sources to a sequence axis, which positionless targets lack), at toy-scale LRs. layers.* expands to +# n_layers sites; the sweep rewrites n_layers (5/50/200) and sites_per_chunk (1 = one +# chunk per layer ... n_layers = 1 chunk == StochasticReconSubsetLoss; must divide +# n_layers). Runs on 1 GPU or CPU. +run_name: jax-deep-linear-30-L5 +pd: + seed: 0 + ci_config: + type: layerwise_mlp + hidden_dims: + - 32 + decomposition_targets: + - module_pattern: layers.* + C: 60 + loss_metrics: + - type: ImportanceMinimalityLoss + coeff: 2e-4 + eps: 1e-6 + pnorm: + start_val: 2.0 + fn_type: linear + final_val_frac: 0.2 + frequency: + coeff: 1e-4 + reference_token_count: 1024 + - type: ChunkwiseSubsetReconLoss + coeff: 1.0 + sites_per_chunk: 1 + routing: + type: uniform_k_subset + - type: FaithfulnessLoss + coeff: 1.0 + batch_size: 1024 + steps: 20000 + components_optimizer: + grad_clip_norm: 0.01 + lr_schedule: + start_val: 1e-3 + fn_type: cosine + warmup_pct: 0.0 + final_val_frac: 0.1 + ci_fn_optimizer: + lr_schedule: + start_val: 1e-3 + fn_type: cosine + warmup_pct: 0.0 + final_val_frac: 0.1 + faithfulness_warmup_steps: 200 + faithfulness_warmup_lr: 0.01 + faithfulness_warmup_weight_decay: 0.0 +runtime: + remat_recon_forwards: false + dp: null +cadence: + train_log_every: 200 + save_every: 5000 + keep_last_n_checkpoints: 2 +target: + n_features: 30 + n_layers: 5 +data: {} +wandb: + project: param-decomp diff --git a/param_decomp_lab/experiments/deep_linear/model.py b/param_decomp_lab/experiments/deep_linear/model.py new file mode 100644 index 000000000..0687bfc0b --- /dev/null +++ b/param_decomp_lab/experiments/deep_linear/model.py @@ -0,0 +1,321 @@ +"""Deep-linear identity target — one-hot inputs through `n_layers` FROZEN identity +layers to a softmax readout; positionless (`leading_axes=()`, the waist is +`[B, n_features]`). + +The target is `logits = x @ I.T @ ... @ I.T` — every site `layers.{i}` is hardcoded to +`eye(n_features)`, nothing is pretrained. The recon comparison is the LM's +`kl_per_position` on the final logits (the softmax lives inside the KL): with one-hot +inputs the clean logit vector IS the input, so the decomposition's ground truth at every +site is the identity — C live rank-1 components `e_i e_iᵀ`, one per feature, with +one-hot per-feature CI (the TMS `identity_ci_error` pattern). +""" + +from dataclasses import dataclass +from typing import Literal + +import equinox as eqx +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P +from jaxtyping import Array, Float + +from param_decomp.components import DecompVU, SiteC, SiteSpec, site_out +from param_decomp.lm import run_stochastic_masked_output +from param_decomp.losses import kl_per_position + +LAYER_PREFIX = "layers." + + +def layer_name(i: int) -> str: + return f"{LAYER_PREFIX}{i}" + + +def site_names_for(n_layers: int) -> tuple[str, ...]: + """Canonical computation order: `layers.0` … `layers.{n_layers-1}`.""" + return tuple(layer_name(i) for i in range(n_layers)) + + +def _parse_layer_index(name: str) -> int: + assert name.startswith(LAYER_PREFIX), f"unknown deep-linear site {name!r}" + return int(name[len(LAYER_PREFIX) :]) + + +ReconKind = Literal["kl", "mse"] + + +@dataclass(frozen=True) +class DeepLinearConfig: + """`logit_scale` multiplies the FINAL output only (a softmax temperature for the KL + recon — the layers stay exact identities, so the per-site ground truth is unchanged); + `recon` picks the recon comparison (`kl` = LM semantics, `mse` = raw logits).""" + + n_features: int + n_layers: int + logit_scale: float = 1.0 + recon: ReconKind = "kl" + + +@dataclass(frozen=True) +class DeepLinearTargetConfig: + """The lab deep-linear target config carried on `ExperimentConfig.target` (satisfies + the core `TargetSites` protocol via `sites`). The target is CONSTRUCTED (identity + weights), not pretrained — a run is reproducible from the config alone. `global_batch` + is the PD-step batch (no parquet `DataConfig`; the toy samples one-hot rows on the + fly).""" + + n_features: int + n_layers: int + logit_scale: float + recon: ReconKind + sites: tuple[SiteC, ...] + global_batch: int + + +class DeepLinearTarget(eqx.Module): + """Frozen identity weights, right-mult oriented like every other target + (`site_out = x @ W.T`; the orientation is moot for `eye`).""" + + layers: tuple[Float[Array, "n_features n_features"], ...] + + +def init_deep_linear_target(cfg: DeepLinearConfig) -> DeepLinearTarget: + return DeepLinearTarget(layers=tuple(jnp.eye(cfg.n_features) for _ in range(cfg.n_layers))) + + +def canonical_site_cs(site_cs: tuple[SiteC, ...]) -> tuple[SiteC, ...]: + """Canonical order is `layers.{i}` in index order (= computation order). The layer + count is inferred from the configured names; each must appear exactly once.""" + by_name = {s.name: s for s in site_cs} + assert len(by_name) == len(site_cs), f"duplicate site in {[s.name for s in site_cs]}" + expected = site_names_for(len(site_cs)) + assert sorted(by_name) == sorted(expected), ( + f"deep-linear sites must be {expected}, got {sorted(by_name)}" + ) + return tuple(by_name[name] for name in expected) + + +def expand_wildcard_site_cs(entries: tuple[SiteC, ...], n_layers: int) -> tuple[SiteC, ...]: + """A single `layers.*` entry expands to every layer at that entry's C (the + `module_pattern` wildcard convention — a 200-layer stack is one config stanza); + explicit `layers.{i}` entries pass through. Result is canonical-ordered.""" + expanded: list[SiteC] = [] + for entry in entries: + if entry.name == f"{LAYER_PREFIX}*": + expanded.extend(SiteC(layer_name(i), entry.C) for i in range(n_layers)) + else: + _parse_layer_index(entry.name) + expanded.append(entry) + assert len(expanded) == n_layers, f"{len(expanded)} expanded sites but n_layers={n_layers}" + return canonical_site_cs(tuple(expanded)) + + +def site_specs(cfg: DeepLinearConfig, site_cs: tuple[SiteC, ...]) -> tuple[SiteSpec, ...]: + site_cs = canonical_site_cs(site_cs) + assert len(site_cs) == cfg.n_layers, f"config n_layers={cfg.n_layers} but {len(site_cs)} sites" + return tuple(SiteSpec(s.name, cfg.n_features, cfg.n_features, s.C) for s in site_cs) + + +def clean_output(target: DeepLinearTarget, resid: Float[Array, "B n_features"]) -> Array: + """The all-frozen forward — the final LOGITS (the softmax lives in `kl_per_position`).""" + for W in target.layers: + resid = resid @ W.T + return resid + + +def site_inputs(target: DeepLinearTarget, resid: Float[Array, "B n_features"]) -> dict[str, Array]: + """Clean CI inputs per site (SPEC S4): each `layers.{i}` reads the frozen output of + the chain up to it.""" + inputs: dict[str, Array] = {} + for i, W in enumerate(target.layers): + inputs[layer_name(i)] = resid + resid = resid @ W.T + return inputs + + +def _run_masked( + target: DeepLinearTarget, + components: DecompVU, + resid: Float[Array, "B n_features"], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + collect: dict[str, Array] | None, +) -> Array: + """The masked decomposed forward (SPEC §4.1, S2): sites in `live` run their decomposed + forward; the rest run the frozen `x @ W.T` path. Each layer reads the (possibly + masked) output of the layer before it, so a masked site propagates.""" + live_set = frozenset(live) + for i, W in enumerate(target.layers): + site = layer_name(i) + if site not in live_set: + resid = resid @ W.T + continue + V, U = components.site(site) + resid = site_out( + resid, V, U, W, masks[site], delta_masks[site] if has_delta else None, + None if routes is None else routes[site], + ) # fmt: skip + if collect is not None: + collect[site] = resid + return resid + + +def weight_deltas_fp32( + target: DeepLinearTarget, components: DecompVU, sites: tuple[SiteSpec, ...] +) -> dict[str, Array]: + """fp32 `W − V@U` per site from fp32 masters (SPEC N2; faithfulness input).""" + out: dict[str, Array] = {} + for spec in sites: + W = target.layers[_parse_layer_index(spec.name)] + V, U = components.site(spec.name) + out[spec.name] = W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T + return out + + +class DeepLinearDecomposedModel(eqx.Module): + """The deep-linear `DecomposedModel` (the `lm.py` contract; SPEC §1), positionless + (`leading_axes=()`). + + Carries the FROZEN `DeepLinearTarget` weights as a field — threaded into the jitted + step as a pytree arg, weights traced not baked. The TRAINABLE V/U (`vu: DecompVU`) is + an explicit method arg, NOT a field (separate lifecycle).""" + + target: DeepLinearTarget + sites: tuple[SiteSpec, ...] = eqx.field(static=True) + leading_axes: tuple[str, ...] = eqx.field(static=True) + logit_scale: float = eqx.field(static=True) + recon: ReconKind = eqx.field(static=True) + + @property + def site_names(self) -> tuple[str, ...]: + return tuple(s.name for s in self.sites) + + def shardings(self, mesh: Mesh) -> "DeepLinearDecomposedModel": + """Replicate every frozen leaf on the mesh — the identity weights are tiny.""" + repl = NamedSharding(mesh, P()) + return jax.tree.map(lambda _a: repl, self) + + def recon_loss_fn( + self, + masked_output: Float[Array, "B n_features"], + clean_output: Float[Array, "B n_features"], + ) -> Array: + """Pure over static config only (no arrays closed over) — safe for the jitted step.""" + if self.recon == "kl": + return kl_per_position(masked_output, clean_output) + masked_output = masked_output.astype(jnp.float32) + clean_output = clean_output.astype(jnp.float32) + return jnp.mean((masked_output - clean_output) ** 2) + + def clean_output(self, resid: Float[Array, "B n_features"]) -> Array: + return clean_output(self.target, resid) * self.logit_scale + + def read_activations( + self, resid: Float[Array, "B n_features"], wanted: tuple[str, ...] + ) -> dict[str, Array]: + inputs = site_inputs(self.target, resid) + return {k: inputs[k] for k in wanted} + + def prepare_compute_weights(self, vu: DecompVU) -> DecompVU: + """Identity: the weights are tiny + replicated, nothing to stack/gather/share.""" + return vu + + def masked_output( + self, + prepared: DecompVU, + resid: Float[Array, "B n_features"], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + *, + remat: bool, + ) -> Array: + def forward( + vu: DecompVU, + resid: Array, + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + ) -> Array: + return ( + _run_masked( + self.target, vu, resid, masks, delta_masks, routes, live, has_delta, None + ) + * self.logit_scale + ) + + forward = jax.checkpoint(forward) if remat else forward + return forward(prepared, resid, masks, delta_masks, routes) + + def stack_ci(self, ci_lower: dict[str, Array]) -> dict[str, Array]: + return ci_lower + + def masked_output_stochastic( + self, + prepared: DecompVU, + resid: Float[Array, "B n_features"], + ci_stacked: dict[str, Array], + draw_key: Array, + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + *, + remat: bool, + ) -> Array: + return run_stochastic_masked_output( + self, prepared, resid, ci_stacked, draw_key, routes, live, has_delta, remat=remat + ) + + def masked_site_outputs( + self, + prepared: DecompVU, + resid: Float[Array, "B n_features"], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + ) -> dict[str, Array]: + collect: dict[str, Array] = {} + _run_masked( + self.target, prepared, resid, masks, delta_masks, routes, live, has_delta, collect + ) + assert set(collect) == set(live), (sorted(collect), sorted(live)) + return collect + + def weight_deltas(self, vu: DecompVU) -> dict[str, Array]: + return weight_deltas_fp32(self.target, vu, self.sites) + + +def deep_linear_decomposed_model( + cfg: DeepLinearConfig, target: DeepLinearTarget, sites: tuple[SiteC, ...] +) -> DeepLinearDecomposedModel: + return DeepLinearDecomposedModel( + target=target, + sites=site_specs(cfg, sites), + leading_axes=(), + logit_scale=cfg.logit_scale, + recon=cfg.recon, + ) + + +def replicate_target[T: (DeepLinearTarget, DeepLinearDecomposedModel)](target: T, mesh: Mesh) -> T: + """Replicate the tiny frozen identity weights on every device.""" + repl = NamedSharding(mesh, P()) + return jax.tree.map(lambda a: jax.device_put(a, repl) if eqx.is_array(a) else a, target) + + +def sample_one_hot(key: Array, batch: int, n_features: int) -> Float[Array, "B n_features"]: + """Uniform one-hot rows — the toy's whole data distribution.""" + return jax.nn.one_hot(jax.random.randint(key, (batch,), 0, n_features), n_features) + + +def one_hot_probe(n_features: int) -> Float[Array, "n_features n_features"]: + """The per-feature CI probe: every one-hot input, magnitude 1.0 (the actual data + values — unlike TMS's 0.75, deep-linear inputs are exactly one-hot).""" + return jnp.eye(n_features) diff --git a/param_decomp_lab/experiments/deep_linear/run.py b/param_decomp_lab/experiments/deep_linear/run.py new file mode 100644 index 000000000..dad15962d --- /dev/null +++ b/param_decomp_lab/experiments/deep_linear/run.py @@ -0,0 +1,175 @@ +"""`pd-deep-linear`: decompose a deep-linear identity model in-process (CPU or 1 GPU). + +Mirrors `pd-tms`, minus pretraining: the target is CONSTRUCTED (`n_layers` frozen +`eye(n_features)` sites), the data is uniform one-hot rows sampled on the fly, and the +recon comparison is the LM's KL-on-final-logits — which makes this the smallest target +that exercises `ChunkwiseSubsetReconLoss` end-to-end. Validated via the ground-truth +identity-CI metric (the TMS `identity_ci_error`, probe magnitude 1.0) every +train-log step. +""" + +from pathlib import Path +from typing import Any + +import equinox as eqx +import fire +import jax +import yaml +from jax import random +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P + +from param_decomp.built_run import LAUNCH_CONFIG_FILENAME, BuiltRun +from param_decomp.components import SiteC +from param_decomp.log import setup_logger +from param_decomp.recon import build_loss_terms +from param_decomp.run import run_decomposition_training +from param_decomp.sharding import hsdp_mesh +from param_decomp.train import TrainState +from param_decomp_lab.experiments import toy_uv_eval +from param_decomp_lab.experiments.config import ( + assert_canonical_algorithm_config, + ci_arch, + run_instance, +) +from param_decomp_lab.experiments.deep_linear import model as deep_linear +from param_decomp_lab.experiments.deep_linear.config import DeepLinearExperimentConfig +from param_decomp_lab.experiments.tms.model import identity_ci_error +from param_decomp_lab.infra.run_files import generate_run_id + + +def build_deep_linear_built_run(cfg: DeepLinearExperimentConfig, run_id: str) -> BuiltRun: + """Convert the canonical schema to the engine's `BuiltRun` bundle via the shared + helpers. Like the other toys, validation is the in-loop target-CI metric, so the + core `eval` stays `None`.""" + site_cs = deep_linear.expand_wildcard_site_cs( + tuple(SiteC(t.module_pattern, t.C) for t in cfg.pd.decomposition_targets), + cfg.target.n_layers, + ) + assert not any(m.type == "PersistentPGDReconLoss" for m in cfg.pd.loss_metrics), ( + "persistent-PGD sources are pinned to a sequence axis in the engine" + " (init_train_state); positionless deep-linear runs cannot carry the adversary" + ) + assert_canonical_algorithm_config(cfg) + build_loss_terms( + cfg.pd.loss_metrics, + tuple(sc.name for sc in site_cs), + cfg.pd.steps, + ) + target = deep_linear.DeepLinearTargetConfig( + n_features=cfg.target.n_features, + n_layers=cfg.target.n_layers, + logit_scale=cfg.target.logit_scale, + recon=cfg.target.recon, + sites=site_cs, + global_batch=cfg.pd.batch_size, + ) + return BuiltRun( + pd=cfg.pd, + runtime=cfg.runtime, + cadence=cfg.cadence, + run=run_instance(cfg, run_id), + target=target, + data=None, + ci_fn=ci_arch(cfg.pd.ci_config, resolve_chunkwise=None), + eval=None, + ) + + +def run_deep_linear_decomposition(built: BuiltRun, raw_cfg: dict[str, Any], mesh: Mesh) -> None: + """Build the frozen identity target, then decompose it through the generic engine.""" + target_cfg = built.target + assert isinstance(target_cfg, deep_linear.DeepLinearTargetConfig) + + dl_cfg = deep_linear.DeepLinearConfig( + n_features=target_cfg.n_features, + n_layers=target_cfg.n_layers, + logit_scale=target_cfg.logit_scale, + recon=target_cfg.recon, + ) + target = deep_linear.init_deep_linear_target(dl_cfg) + lm = deep_linear.replicate_target( + deep_linear.deep_linear_decomposed_model(dl_cfg, target, target_cfg.sites), mesh + ) + + data_key = random.fold_in(random.PRNGKey(built.pd.seed), 17) + + @jax.jit + def sample_one_hot_batch(step_key: jax.Array) -> jax.Array: + x = deep_linear.sample_one_hot(step_key, target_cfg.global_batch, target_cfg.n_features) + return jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P(("replicate", "fsdp")))) + + def sample_batch(step: int) -> jax.Array: + return sample_one_hot_batch(random.fold_in(data_key, step)) + + # `model` is the filter_jit ARG (frozen weights traced, not baked) — closing over an + # array-bearing eqx model would bake its weights into the HLO. + @eqx.filter_jit + def probe_ci( + model: deep_linear.DeepLinearDecomposedModel, ci_fn: Any + ) -> tuple[dict[str, jax.Array], dict[str, jax.Array]]: + probe = deep_linear.one_hot_probe(target_cfg.n_features) + ci = ci_fn(model.read_activations(probe, ci_fn.input_names), remat=False) + return ci.lower, ci.upper + + uv_spec = toy_uv_eval.toy_uv_spec(lm, raw_cfg) + + def eval_fn(state: TrainState, now_step: int) -> dict[str, float]: + ci_lower, ci_upper = probe_ci(lm, state.ci_fn) + toy_uv_eval.log_uv_figure( + uv_spec, + state.components.vu, + ci_upper, + now_step, + wandb_active=built.run.wandb is not None, + ) + return { + f"eval/identity_ci_error/{site}": float(identity_ci_error(ci, tolerance=0.1)) + for site, ci in ci_lower.items() + } + + run_decomposition_training( + pd=built.pd, + cadence=built.cadence, + run=built.run, + lm=lm, + ci_fn=built.ci_fn, + data=built.data, + remat_recon_forwards=built.runtime.remat_recon_forwards, + remat_ci_fn=built.runtime.remat_ci_fn, + ascend_replicate=built.runtime.ascend_replicate, + compiler_options=built.runtime.compiler_options, + profile=built.runtime.launch_env.profile, + sample_batch=sample_batch, + eval_fn=eval_fn, + eval_every=built.cadence.train_log_every, + mesh=mesh, + ) + + +def main(config: str, group: str | None = None, tags: str | None = None) -> None: + schema_raw = yaml.safe_load(Path(config).read_text()) + run_id = generate_run_id("param_decomp") + if group is not None or tags is not None: + wandb_cfg = dict(schema_raw.get("wandb") or {}) + if group is not None: + wandb_cfg["group"] = group + if tags is not None: + wandb_cfg["tags"] = tags.split(",") + schema_raw["wandb"] = wandb_cfg + built = build_deep_linear_built_run(DeepLinearExperimentConfig(**schema_raw), run_id) + built.run.run_dir.mkdir(parents=True, exist_ok=True) + setup_logger(built.run.run_dir / "logs.log") + (built.run.run_dir / LAUNCH_CONFIG_FILENAME).write_text( + yaml.safe_dump(schema_raw, sort_keys=False) + ) + mesh = hsdp_mesh() + run_deep_linear_decomposition(built, schema_raw, mesh) + + +def cli() -> None: + fire.Fire(main) + + +if __name__ == "__main__": + cli() diff --git a/param_decomp_lab/experiments/deep_linear/test_deep_linear.py b/param_decomp_lab/experiments/deep_linear/test_deep_linear.py new file mode 100644 index 000000000..ce8c66252 --- /dev/null +++ b/param_decomp_lab/experiments/deep_linear/test_deep_linear.py @@ -0,0 +1,254 @@ +"""CPU tests for the deep-linear identity target over the generic positionless core. + +Covers the `DecomposedModel` contract (clean == frozen passthrough of one-hot inputs, +masked identity, KL recon), chunkwise plan construction over the homogeneous sites, the +full SPEC step with `ChunkwiseSubsetReconLoss` (the first positionless target to +exercise it), the loud rejection of the seq-axis-pinned persistent-PGD adversary, and an +end-to-end decompose → recovers-identity validation on a tiny d=6 stack. +""" + +from pathlib import Path + +import equinox as eqx +import jax +import jax.numpy as jnp +import optax +import pytest +import yaml + +from param_decomp.ci_fn import MLPCIArch, init_layerwise_mlp_ci_fn +from param_decomp.components import SiteC, SiteSpec, init_decomp_vu +from param_decomp.configs import ( + ChunkwiseSubsetReconLossConfig, + FaithfulnessLossConfig, + ImportanceMinimalityLossConfig, +) +from param_decomp.losses import kl_per_position +from param_decomp.recon import ReconLossTerm, build_loss_terms +from param_decomp.schedule import ScheduleConfig +from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step +from param_decomp_lab.experiments.deep_linear.config import DeepLinearExperimentConfig +from param_decomp_lab.experiments.deep_linear.model import ( + DeepLinearConfig, + canonical_site_cs, + deep_linear_decomposed_model, + expand_wildcard_site_cs, + init_deep_linear_target, + one_hot_probe, + sample_one_hot, + site_inputs, + site_specs, +) +from param_decomp_lab.experiments.deep_linear.run import build_deep_linear_built_run +from param_decomp_lab.experiments.tms.model import identity_ci_error + +CONFIGS_DIR = Path(__file__).parent / "configs" + + +def _cfg(n_features: int = 6, n_layers: int = 4) -> DeepLinearConfig: + return DeepLinearConfig(n_features=n_features, n_layers=n_layers) + + +def _site_cs(n_layers: int = 4, c: int = 12) -> tuple[SiteC, ...]: + return tuple(SiteC(f"layers.{i}", c) for i in range(n_layers)) + + +def test_canonical_order_and_dims(): + cfg = _cfg() + shuffled = (_site_cs()[2], _site_cs()[0], _site_cs()[3], _site_cs()[1]) + assert canonical_site_cs(shuffled) == _site_cs() + with pytest.raises(AssertionError): + canonical_site_cs(_site_cs()[:2] + _site_cs()[:1]) # duplicate site + specs = site_specs(cfg, _site_cs()) + assert all((s.d_in, s.d_out, s.C) == (6, 6, 12) for s in specs) + + +def test_wildcard_expansion(): + assert expand_wildcard_site_cs((SiteC("layers.*", 12),), 4) == _site_cs() + # explicit names pass through; count must match n_layers + assert expand_wildcard_site_cs(_site_cs(), 4) == _site_cs() + with pytest.raises(AssertionError): + expand_wildcard_site_cs((SiteC("layers.*", 12),) + _site_cs()[:1], 4) + + +def test_clean_output_is_one_hot_passthrough(): + cfg = _cfg() + target = init_deep_linear_target(cfg) + lm = deep_linear_decomposed_model(cfg, target, _site_cs()) + x = sample_one_hot(jax.random.PRNGKey(0), 16, cfg.n_features) + clean = lm.clean_output(x) + assert jnp.array_equal(clean, x), "identity stack must pass one-hot inputs through" + # every site reads the same one-hot activation + site_in = site_inputs(target, x) + assert all(jnp.array_equal(v, x) for v in site_in.values()) + + +def test_masked_identity_and_ablation(): + cfg = _cfg() + sites = site_specs(cfg, _site_cs()) + target = init_deep_linear_target(cfg) + lm = deep_linear_decomposed_model(cfg, target, _site_cs()) + vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) + b = 7 + x = sample_one_hot(jax.random.PRNGKey(2), b, cfg.n_features) + clean = lm.clean_output(x) + + # SPEC S2: live=() is the exact frozen path. + none_masked = lm.masked_output(vu, x, {}, {}, None, (), True, remat=False) + assert jnp.array_equal(clean, none_masked) + + # All-live, masks=1, delta=1 reconstructs the frozen path up to decomposition rounding. + names = lm.site_names + ones_masks = {s.name: jnp.ones((b, s.C)) for s in lm.sites} + ones_delta = {s: jnp.ones((b,)) for s in names} + full = lm.masked_output(vu, x, ones_masks, ones_delta, None, names, True, remat=False) + assert jnp.allclose(clean, full, atol=1e-4) + + # Zero-masking one site changes the output. + C = lm.sites[0].C + ablated = lm.masked_output( + vu, x, {"layers.0": jnp.zeros((b, C))}, {"layers.0": jnp.zeros((b,))}, + None, ("layers.0",), True, remat=False, + ) # fmt: skip + assert not jnp.allclose(clean, ablated, atol=1e-5) + + deltas = lm.weight_deltas(vu) + assert all(v.shape == (cfg.n_features, cfg.n_features) for v in deltas.values()) + assert all(v.dtype == jnp.float32 for v in deltas.values()) + + +def test_recon_loss_fn_is_kl(): + cfg = _cfg() + lm = deep_linear_decomposed_model(cfg, init_deep_linear_target(cfg), _site_cs()) + a = jax.random.normal(jax.random.PRNGKey(0), (4, cfg.n_features)) + b = jax.random.normal(jax.random.PRNGKey(1), (4, cfg.n_features)) + assert jnp.array_equal(lm.recon_loss_fn(a, b), kl_per_position(a, b)) + assert float(lm.recon_loss_fn(a, a)) == pytest.approx(0.0, abs=1e-6) + assert float(lm.recon_loss_fn(a, b)) > 0.0 + + +@pytest.mark.parametrize(("sites_per_chunk", "n_chunks"), [(1, 4), (2, 2), (4, 1)]) +def test_chunkwise_plan_over_layers(sites_per_chunk: int, n_chunks: int): + names = tuple(f"layers.{i}" for i in range(4)) + terms = build_loss_terms( + ( + FaithfulnessLossConfig(coeff=1.0), + ImportanceMinimalityLossConfig( + coeff=1e-3, pnorm=ScheduleConfig(start_val=1.0, fn_type="constant") + ), + ChunkwiseSubsetReconLossConfig(coeff=1.0, sites_per_chunk=sites_per_chunk), + ), + names, + 100, + ) + (term,) = terms.recon + assert isinstance(term, ReconLossTerm) + assert len(term.plan) == n_chunks + covered = tuple(site for entry in term.plan for site in entry.live_sites) + assert covered == names, "chunks must tile the sites in canonical order" + + +def _loss_metrics(sites_per_chunk: int): + return ( + FaithfulnessLossConfig(coeff=1.0), + ImportanceMinimalityLossConfig( + coeff=1e-3, pnorm=ScheduleConfig(start_val=1.0, fn_type="constant") + ), + ChunkwiseSubsetReconLossConfig(coeff=1.0, sites_per_chunk=sites_per_chunk), + ) + + +def _state_and_step( + cfg: DeepLinearConfig, + sites: tuple[SiteSpec, ...], + sites_per_chunk: int, + total_steps: int, + warmup_steps: int = 0, +): + lm = deep_linear_decomposed_model( + cfg, init_deep_linear_target(cfg), _site_cs(cfg.n_layers, sites[0].C) + ) + vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) + if warmup_steps: + warm_opt = optax.adamw(1e-2, weight_decay=0.0) + wstep = make_faith_warmup_step(warm_opt) + warm_state = warm_opt.init(eqx.filter(vu, eqx.is_array)) + for _ in range(warmup_steps): + vu, warm_state, _ = wstep(lm, vu, warm_state) + ci_fn = init_layerwise_mlp_ci_fn(MLPCIArch(hidden_dims=(16,)), sites, jax.random.PRNGKey(2)) + 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) + loss_terms = build_loss_terms(_loss_metrics(sites_per_chunk), lm.site_names, total_steps) + state = TrainState( + components=vu, ci_fn=ci_fn, + components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)), + ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)), + adversaries={}, step=jnp.zeros((), jnp.int32), + ) # fmt: skip + 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, + ) # fmt: skip + return lm, state, step, loss_terms + + +@pytest.mark.parametrize("sites_per_chunk", [1, 4]) +def test_step_trains_chunkwise_stochastic(sites_per_chunk: int): + cfg = _cfg() + sites = site_specs(cfg, _site_cs()) + lm, state, step, _ = _state_and_step(cfg, sites, sites_per_chunk, total_steps=20) + for i in range(6): + x = sample_one_hot(jax.random.fold_in(jax.random.PRNGKey(99), i), 64, cfg.n_features) + state, m = step(lm, state, x, jax.random.PRNGKey(100 + i)) + assert all(bool(jnp.isfinite(v).all()) for v in m.values()), m + assert int(state.step) == 6 + assert state.adversaries == {} + + +def test_persistent_pgd_rejected_loudly(): + """The engine pins persistent sources to a sequence axis (`init_train_state` requires + a `DataConfig`); the deep-linear build must reject the adversary up front, not at + engine init.""" + raw = yaml.safe_load((CONFIGS_DIR / "deep_linear_30.yaml").read_text()) + raw["pd"]["loss_metrics"].append( + { + "type": "PersistentPGDReconLoss", + "coeff": 0.5, + "optimizer": { + "type": "adam", + "lr_schedule": {"start_val": 0.01, "fn_type": "constant"}, + }, + "scope": {"type": "bsc"}, + } + ) + with pytest.raises(AssertionError, match="sequence axis"): + build_deep_linear_built_run(DeepLinearExperimentConfig(**raw), "p-00000000") + + +@pytest.mark.parametrize("config_name", ["deep_linear_30.yaml", "deep_linear_128.yaml"]) +def test_repo_configs_parse_and_build(config_name: str): + raw = yaml.safe_load((CONFIGS_DIR / config_name).read_text()) + built = build_deep_linear_built_run(DeepLinearExperimentConfig(**raw), "p-00000000") + assert built.data is None and built.eval is None + assert len(built.target.sites) == 5 + + +@pytest.mark.slow +def test_end_to_end_decompose_recovers_identity(): + """The recovers-structure gate: decompose a d=6, 2-layer identity stack with the + chunkwise stochastic recon (one chunk per layer) and show the recovered per-site CI + of the one-hot probe is the identity up to permutation — zero `identity_ci_error`.""" + cfg = _cfg(n_features=6, n_layers=2) + site_cs = _site_cs(n_layers=2, c=12) + sites = site_specs(cfg, site_cs) + lm, state, step, _ = _state_and_step( + cfg, sites, sites_per_chunk=1, total_steps=3000, warmup_steps=200 + ) + for i in range(3000): + x = sample_one_hot(jax.random.fold_in(jax.random.PRNGKey(5), i), 512, cfg.n_features) + state, _ = step(lm, state, x, jax.random.fold_in(jax.random.PRNGKey(6), i)) + + probe = one_hot_probe(cfg.n_features) + ci = state.ci_fn(lm.read_activations(probe, state.ci_fn.input_names), remat=False) + errors = {site: identity_ci_error(v, tolerance=0.2) for site, v in ci.lower.items()} + assert all(e == 0 for e in errors.values()), errors diff --git a/param_decomp_lab/experiments/lm/config.py b/param_decomp_lab/experiments/lm/config.py index f49d3062f..d1c436c3b 100644 --- a/param_decomp_lab/experiments/lm/config.py +++ b/param_decomp_lab/experiments/lm/config.py @@ -308,7 +308,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..c7969496f 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, @@ -149,7 +150,7 @@ def single_feature_ci( model: resid_mlp.ResidMLPDecomposedModel, ci_fn: Any ) -> tuple[dict[str, jax.Array], dict[str, jax.Array]]: resid = resid_mlp.single_feature_probe(target_cfg.n_features) @ model.target.W_E - ci = ci_fn(model.read_activations(resid, ci_fn.input_names)) + ci = ci_fn(model.read_activations(resid, ci_fn.input_names), remat=False) return ci.lower, ci.upper uv_spec = toy_uv_eval.toy_uv_spec(lm, raw_cfg) 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..9c6037576 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, @@ -132,7 +133,7 @@ def single_feature_ci( model: tms.TMSDecomposedModel, ci_fn: Any ) -> tuple[dict[str, jax.Array], dict[str, jax.Array]]: probe = tms.single_feature_probe(target_cfg.n_features) - ci = ci_fn(model.read_activations(probe, ci_fn.input_names)) + ci = ci_fn(model.read_activations(probe, ci_fn.input_names), remat=False) return ci.lower, ci.upper uv_spec = toy_uv_eval.toy_uv_spec(lm, raw_cfg) 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, diff --git a/param_decomp_lab/pyproject.toml b/param_decomp_lab/pyproject.toml index cd16cd26d..e6a20c691 100644 --- a/param_decomp_lab/pyproject.toml +++ b/param_decomp_lab/pyproject.toml @@ -41,6 +41,7 @@ pd-lm = "param_decomp_lab.experiments.lm.launch:cli" # core engine `run_decomposition_training`). Run synchronously, no SLURM. pd-tms = "param_decomp_lab.experiments.tms.run:cli" pd-resid-mlp = "param_decomp_lab.experiments.resid_mlp.run:cli" +pd-deep-linear = "param_decomp_lab.experiments.deep_linear.run:cli" # Target-model pretraining launcher: snapshot + workspace + sbatch # `python -m pretrain.train` (or `--local`). The core pretrain trainer lives at the root # sibling `pretrain/`.