fix(consumers): metadata consumers parse only the config subset they read

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:

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)
 

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)

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,29 +229,25 @@ 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()
             return 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],
             )

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)