From bf8c4f12491a7280d32f88290e51633c3edc32a9 Mon Sep 17 00:00:00 2001 From: oli Date: Mon, 6 Jul 2026 10:25:15 +0000 Subject: [PATCH] fix(consumers): metadata consumers parse only the config subset they read run_metadata / PDAdapter re-validated a run's pinned launch_config.yaml against the full LMExperimentConfig (extra=forbid), so every trainer-schema change locked consumers out of stored runs: at tip only 7/104 stored runs parsed (post-#939 anneal fields, removed data/runtime/optimizer fields, deleted loss-type union tags, and fields from branches ahead of tip). ConsumerRunConfig parses only what those consumers read (target.spec, pd.decomposition_targets, four data fields) with extra=ignore, strict on the read fields. Target resolution is factored into resolve_target so it no longer needs the full schema. run_metadata now opens 103/103 stored runs. open_jax_run keeps the full-schema load_run_dir_config: it must rebuild the exact TrainState pytree (optimizers + per-loss-term persistent state) to restore orbax. Shrinking it the same way is gated on the checkpoint-split work. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W9D6bgyT239nJ3E9gF3S3s --- param_decomp_lab/adapters/pd.py | 6 +- param_decomp_lab/experiments/CLAUDE.md | 6 +- param_decomp_lab/experiments/lm/config.py | 101 ++++++++++++++---- param_decomp_lab/experiments/lm/load_run.py | 21 ++-- .../tests/experiments/test_consumer_config.py | 75 +++++++++++++ 5 files changed, 174 insertions(+), 35 deletions(-) create mode 100644 param_decomp_lab/tests/experiments/test_consumer_config.py diff --git a/param_decomp_lab/adapters/pd.py b/param_decomp_lab/adapters/pd.py index c69501f35..a045eb1de 100644 --- a/param_decomp_lab/adapters/pd.py +++ b/param_decomp_lab/adapters/pd.py @@ -3,7 +3,7 @@ from param_decomp.built_run import LAUNCH_CONFIG_FILENAME from param_decomp_lab.autointerp.schemas import ModelMetadata -from param_decomp_lab.experiments.lm.config import LMExperimentConfig +from param_decomp_lab.experiments.lm.config import ConsumerRunConfig, load_consumer_config from param_decomp_lab.experiments.lm.load_run import RunMetadata, run_metadata from param_decomp_lab.harvest.schemas import get_harvest_dir from param_decomp_lab.topology.path_schemas import path_schema_for_model_type @@ -32,10 +32,10 @@ def _run_dir(self) -> Path: return get_harvest_dir(self._run_id).parent @cached_property - def cfg(self) -> LMExperimentConfig: + def cfg(self) -> ConsumerRunConfig: config_path = self._run_dir / LAUNCH_CONFIG_FILENAME assert config_path.exists(), f"config not found: {config_path}" - return LMExperimentConfig.from_file(config_path) + return load_consumer_config(self._run_dir) @cached_property def _metadata(self) -> RunMetadata: diff --git a/param_decomp_lab/experiments/CLAUDE.md b/param_decomp_lab/experiments/CLAUDE.md index 47a8ea4cd..4e0a9d321 100644 --- a/param_decomp_lab/experiments/CLAUDE.md +++ b/param_decomp_lab/experiments/CLAUDE.md @@ -12,7 +12,11 @@ validation / run-identity helpers (`assert_canonical_algorithm_config` / `run_in target/data schema + (for the LM) its `BuiltRun` build. 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`. +see `param_decomp_lab/adapters/pd.py`. That path parses only the consumer-read subset of +the pinned config (`ConsumerRunConfig`, `extra="ignore"`), so trainer-schema drift in +either direction doesn't lock consumers out of stored runs; `open_jax_run` still +validates the full `LMExperimentConfig` because restoring orbax needs the exact +`TrainState` pytree. ## Toy domains (TMS, ResidMLP) diff --git a/param_decomp_lab/experiments/lm/config.py b/param_decomp_lab/experiments/lm/config.py index b70313ce6..744fcead4 100644 --- a/param_decomp_lab/experiments/lm/config.py +++ b/param_decomp_lab/experiments/lm/config.py @@ -5,15 +5,17 @@ `BuiltRun` bundle (`param_decomp.built_run`) — the pydantic `pd` / `cadence` / `runtime` verbatim plus the resolved target / data / CI-fn arch / eval — asserting loudly on anything the JAX trainer doesn't implement. The composition entry (`run.py`) calls `load_config` / -`build_from_schema`; consumers that read a finished run dir call `load_run_dir_config`. +`build_from_schema`; `open_jax_run` (which must rebuild the full `TrainState`) calls +`load_run_dir_config`; the metadata consumers (`run_metadata` / `PDAdapter`) call +`load_consumer_config`, a drift-tolerant parse of only the fields they read. """ from dataclasses import dataclass from pathlib import Path -from typing import Annotated, Any, Literal +from typing import Annotated, Any, ClassVar, Literal import yaml -from pydantic import Discriminator, Field, PositiveInt, model_validator +from pydantic import ConfigDict, Discriminator, Field, PositiveInt, model_validator from param_decomp.base_config import BaseConfig from param_decomp.built_run import ( @@ -200,29 +202,29 @@ class LlamaSimpleMLPTargetConfig: ) -def _site_cs(cfg: LMExperimentConfig) -> tuple[SiteC, ...]: +def _canonical_llama_site_cs(raw_site_cs: tuple[SiteC, ...]) -> tuple[SiteC, ...]: """Decomposition targets -> canonical per-site (name, C) pairs. Any per-layer matrix site (q/k/v/o/gate/up/down) with its own C is supported; non-site module patterns refuse. Raw-HF specs name modules `model.layers.*`; the vendored class drops the prefix — same matrices either way.""" site_cs = [] - for target in cfg.pd.decomposition_targets: - name = target.module_pattern.removeprefix("model.") - assert SITE_NAME_PATTERN.match(name), ( - f"unsupported decomposition target {target.module_pattern!r}" - ) - site_cs.append(SiteC(name, target.C)) + for raw in raw_site_cs: + name = raw.name.removeprefix("model.") + assert SITE_NAME_PATTERN.match(name), f"unsupported decomposition target {raw.name!r}" + site_cs.append(SiteC(name, raw.C)) return canonical_site_cs(tuple(site_cs)) -def _resolve_target(cfg: LMExperimentConfig) -> AnyLMTargetConfig: - """Target spec + decomposition patterns -> the JAX target config. +def resolve_target( + spec: LMTargetSpec, raw_site_cs: tuple[SiteC, ...], max_seq_len: int +) -> AnyLMTargetConfig: + """Target spec + decomposition patterns (module_pattern/C as configured) -> the JAX + target config. Vendored and raw-HF Llama specs load the SAME meta-llama weights (the export bridge round-trip verified vendored == HF numerics); both map to the HF loader. `kind: pretrained` LlamaSimpleMLP specs map to the pretrain-cache loader, with `h.*` wildcard patterns expanded over the checkpoint's n_layer.""" - spec = cfg.target.spec match spec: case HFWeightsInVendored() | HFTarget(): match spec: @@ -231,19 +233,26 @@ def _resolve_target(cfg: LMExperimentConfig) -> AnyLMTargetConfig: case HFTarget(): assert spec.model_class == "transformers.LlamaForCausalLM", spec.model_class assert "Llama-3.1-8B" in spec.model_name, spec.model_name - return TargetConfig(model_name=spec.model_name, sites=_site_cs(cfg)) + return TargetConfig( + model_name=spec.model_name, sites=_canonical_llama_site_cs(raw_site_cs) + ) case PretrainedTarget(): assert spec.model_class.rsplit(".", 1)[-1] == "LlamaSimpleMLP", spec.model_class cache_dir = llama_simple_mlp.pretrain_cache_dir(spec.run_path) arch = llama_simple_mlp.load_model_config(cache_dir) - assert cfg.data.max_seq_len <= arch.n_ctx, (cfg.data.max_seq_len, arch.n_ctx) - sites = llama_simple_mlp.expand_wildcard_site_cs( - tuple(SiteC(t.module_pattern, t.C) for t in cfg.pd.decomposition_targets), - arch.n_layer, - ) + assert max_seq_len <= arch.n_ctx, (max_seq_len, arch.n_ctx) + sites = llama_simple_mlp.expand_wildcard_site_cs(raw_site_cs, arch.n_layer) return LlamaSimpleMLPTargetConfig(pretrain_run_path=spec.run_path, sites=sites) +def _resolve_target(cfg: LMExperimentConfig) -> AnyLMTargetConfig: + return resolve_target( + cfg.target.spec, + tuple(SiteC(t.module_pattern, t.C) for t in cfg.pd.decomposition_targets), + cfg.data.max_seq_len, + ) + + def _block_of_site(target: AnyLMTargetConfig, site_name: str) -> int: """Transformer-block index of a decomposition site, parsed with the target's own site grammar (`layers.{i}...` for llama8b, `h.{i}...` for LlamaSimpleMLP).""" @@ -460,3 +469,57 @@ def load_run_dir_config(run_dir: Path) -> BuiltRun: run id is the run-dir name.""" schema_raw = yaml.safe_load((run_dir / LAUNCH_CONFIG_FILENAME).read_text()) return build_from_schema(schema_raw, run_dir.name) + + +class _ConsumerSchema(BaseConfig): + """Base for the consumer-read subset of a pinned launch config: strict on the fields + consumers read, `extra="ignore"` on everything else, so trainer-schema drift — a field + removed at tip, or a run launched from a branch ahead of tip — cannot lock consumers + out of stored runs. The full `LMExperimentConfig` (`extra="forbid"`) remains the + train/resume contract.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="ignore") + + +class ConsumerDecompositionTarget(_ConsumerSchema): + module_pattern: str + C: PositiveInt + + +class ConsumerTargetConfig(_ConsumerSchema): + spec: LMTargetSpec + + +class ConsumerPDConfig(_ConsumerSchema): + decomposition_targets: tuple[ConsumerDecompositionTarget, ...] + + +class ConsumerDataConfig(_ConsumerSchema): + dataset_name: str + data_files: str | None = None + tokenizer_name: str + max_seq_len: PositiveInt + + +class ConsumerRunConfig(_ConsumerSchema): + """What the metadata consumers (`run_metadata` / `PDAdapter`, feeding + autointerp/clustering) read from a run's pinned launch config. `open_jax_run` still + validates the full schema — it must rebuild the exact `TrainState` pytree (optimizers + + per-loss-term persistent state) to restore the orbax checkpoint.""" + + target: ConsumerTargetConfig + pd: ConsumerPDConfig + data: ConsumerDataConfig + + def resolved_target(self) -> AnyLMTargetConfig: + return resolve_target( + self.target.spec, + tuple(SiteC(t.module_pattern, t.C) for t in self.pd.decomposition_targets), + self.data.max_seq_len, + ) + + +def load_consumer_config(run_dir: Path) -> ConsumerRunConfig: + """Parse the consumer-read subset of `run_dir`'s pinned launch config. Tolerates + trainer-schema drift outside the read fields (see `_ConsumerSchema`).""" + return ConsumerRunConfig.from_file(run_dir / LAUNCH_CONFIG_FILENAME) diff --git a/param_decomp_lab/experiments/lm/load_run.py b/param_decomp_lab/experiments/lm/load_run.py index c064d8591..5bb665034 100644 --- a/param_decomp_lab/experiments/lm/load_run.py +++ b/param_decomp_lab/experiments/lm/load_run.py @@ -50,6 +50,7 @@ from param_decomp_lab.experiments.lm.config import ( LlamaSimpleMLPTargetConfig, TargetConfig, + load_consumer_config, load_run_dir_config, ) @@ -228,18 +229,19 @@ class RunMetadata: def run_metadata(run_dir: Path) -> RunMetadata: - """Target topology for `run_dir`, derived from the pinned config (+ the SimpleMLP - pretrain cache's `model_config.yaml` for `n_layer`/`vocab_size`). No orbax restore.""" - cfg = load_run_dir_config(run_dir) - match cfg.target: + """Target topology for `run_dir`, derived from the consumer-read subset of the pinned + config (drift-tolerant, see `load_consumer_config`; + the SimpleMLP pretrain cache's + `model_config.yaml` for `n_layer`/`vocab_size`). No orbax restore.""" + target = load_consumer_config(run_dir).resolved_target() + match target: case LlamaSimpleMLPTargetConfig(): - cache_dir = llama_simple_mlp.pretrain_cache_dir(cfg.target.pretrain_run_path) + cache_dir = llama_simple_mlp.pretrain_cache_dir(target.pretrain_run_path) simple_cfg = llama_simple_mlp.load_model_config(cache_dir) return RunMetadata( model_type="LlamaSimpleMLP", n_blocks=simple_cfg.n_layer, vocab_size=simple_cfg.vocab_size, - layer_activation_sizes=[(s.name, s.C) for s in cfg.target.sites], + layer_activation_sizes=[(s.name, s.C) for s in target.sites], ) case TargetConfig(): llama_cfg = llama31_8b_config() @@ -247,10 +249,5 @@ def run_metadata(run_dir: Path) -> RunMetadata: model_type="Llama", n_blocks=llama_cfg.n_layer, vocab_size=llama_cfg.vocab_size, - layer_activation_sizes=[(s.name, s.C) for s in cfg.target.sites], - ) - case _: - raise AssertionError( - "run_metadata is the LM-consumer path only (toys are not harvested); " - f"got target {type(cfg.target).__name__}" + layer_activation_sizes=[(s.name, s.C) for s in target.sites], ) diff --git a/param_decomp_lab/tests/experiments/test_consumer_config.py b/param_decomp_lab/tests/experiments/test_consumer_config.py new file mode 100644 index 000000000..eba76e518 --- /dev/null +++ b/param_decomp_lab/tests/experiments/test_consumer_config.py @@ -0,0 +1,75 @@ +"""`ConsumerRunConfig` opens stored run configs across trainer-schema drift: strict on the +fields consumers read, indifferent to everything else. The full `LMExperimentConfig` +(`extra="forbid"`) is expected to REFUSE the same drifted config — that contrast is the +contract under test.""" + +from typing import Any + +import pytest +from pydantic import ValidationError + +from param_decomp_lab.experiments.lm.config import ConsumerRunConfig, LMExperimentConfig + +DRIFTED_LAUNCH_CONFIG: dict[str, Any] = { + "run_name": "drifted", + "removed_top_level_knob": True, + "target": { + "spec": { + "kind": "pretrained", + "model_class": "param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP", + "run_path": "goodfire/spd/runs/abcd1234", + }, + "weights_dtype": "bfloat16", + }, + "pd": { + "seed": 0, + "batch_size": 64, + "steps": 1000, + "decomposition_targets": [ + {"module_pattern": "h.*.mlp.c_fc", "C": 1200}, + {"module_pattern": "h.*.mlp.down_proj", "C": 1200}, + ], + "loss_metrics": [{"type": "LossTypeDeletedFromTip", "coeff": 1.0}], + }, + "data": { + "dataset_name": "parquet", + "data_files": "/shards/pile_neox_tok_512/*.parquet", + "tokenizer_name": "EleutherAI/gpt-neox-20b", + "max_seq_len": 512, + "column_name": "input_ids", + "field_from_a_branch_ahead_of_tip": 4, + }, + "runtime": {"dp": 64, "field_removed_at_tip": True}, +} + + +def test_consumer_config_parses_across_drift_where_full_schema_refuses(): + cfg = ConsumerRunConfig.model_validate(DRIFTED_LAUNCH_CONFIG) + assert cfg.target.spec.kind == "pretrained" + assert [(t.module_pattern, t.C) for t in cfg.pd.decomposition_targets] == [ + ("h.*.mlp.c_fc", 1200), + ("h.*.mlp.down_proj", 1200), + ] + assert cfg.data.tokenizer_name == "EleutherAI/gpt-neox-20b" + assert cfg.data.max_seq_len == 512 + + with pytest.raises(ValidationError): + LMExperimentConfig.model_validate(DRIFTED_LAUNCH_CONFIG) + + +def test_consumer_config_is_strict_on_the_fields_it_reads(): + missing_targets = { + **DRIFTED_LAUNCH_CONFIG, + "pd": { + k: v for k, v in DRIFTED_LAUNCH_CONFIG["pd"].items() if k != "decomposition_targets" + }, + } + with pytest.raises(ValidationError): + ConsumerRunConfig.model_validate(missing_targets) + + bad_spec = { + **DRIFTED_LAUNCH_CONFIG, + "target": {"spec": {"kind": "no_such_kind"}}, + } + with pytest.raises(ValidationError): + ConsumerRunConfig.model_validate(bad_spec)