fix(consumers): metadata consumers parse only the config subset they read#952
Open
ocg-goodfire wants to merge 1 commit into
Open
fix(consumers): metadata consumers parse only the config subset they read#952ocg-goodfire wants to merge 1 commit into
ocg-goodfire wants to merge 1 commit into
Conversation
…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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01W9D6bgyT239nJ3E9gF3S3s
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
Stored runs' pinned
launch_config.yamlrefuse to parse under the current schema, so consumers can't open them. At feature/jax tip only 7 of 104 stored runs parse — the #939 anneal fields are the dominant cohort, but the sweep found ~10 distinct drift kinds: removeddata/runtime/optimizer fields, deleted loss-type union tags, and fields from branches ahead of tip (runs launched today from an unmerged branch usedata.train_batch_replay, unknown at tip). Every future schema change recreates this, in both directions.The root cause for the metadata consumers is accidental coupling:
run_metadata/PDAdapter(feeding autointerp/clustering) re-validated the fullLMExperimentConfig(extra="forbid") when all they read istarget.spec,pd.decomposition_targets, and fourdatafields.What
ConsumerRunConfig(+load_consumer_config): parses only the consumer-read subset,extra="ignore"elsewhere, strict on the read fields. The full schema remains the train/resume contract.resolve_target: target resolution factored out of_resolve_target(LMExperimentConfig)so it works from the subset (theLMExperimentConfigpath delegates to it).run_metadataandPDAdapter.cfgswitched onto the consumer parse. Therun_metadatatoy catch-all is gone —resolve_targetreturns the closed LM union, so the match is exhaustive by type.ConsumerRunConfigand refuses viaLMExperimentConfig; missing/invalid read-fields still refuse.Swept against the real runs dir:
run_metadatanow opens 103/103 stored runs (was 7/104 via the full parse; one run dir gained a config mid-sweep, hence 103 vs 104).What this does NOT fix
open_jax_run(harvest/clustering workers) keeps the full-schemaload_run_dir_config: it must rebuild the exactTrainStatepytree (optimizers + per-loss-term persistent state, e.g. PPGD sources) to restore orbax, so it genuinely depends on the full config. Unblocking it for existing runs is the one-shot #939 migration (script attached to that PR, never run — coordinated separately in the bridge thread); shrinking it structurally is gated on the checkpoint-split work (decomposition item = components + ci_fn only).Validation
make checkcleanpytest param_decomp_lab -m 'not slow': 230 passedrun_metadataover all stored runs on cluster: 103/103 OKbridge thread: consumer-config-schema-drift
🤖 Generated with Claude Code
https://claude.ai/code/session_01W9D6bgyT239nJ3E9gF3S3s