Skip to content
This repository was archived by the owner on Jul 17, 2026. It is now read-only.

refactor(config): per-domain decomposition {sites, ci} — carve, migration, and site-grammar cleanup#966

Open
claude-spd1 wants to merge 7 commits into
feature/jaxfrom
bridge/task-targets-refactor
Open

refactor(config): per-domain decomposition {sites, ci} — carve, migration, and site-grammar cleanup#966
claude-spd1 wants to merge 7 commits into
feature/jaxfrom
bridge/task-targets-refactor

Conversation

@claude-spd1

Copy link
Copy Markdown
Collaborator

Lands the decomposition-config carve (Oli's refactor/decomposition-config-shape, rebased onto feature/jax tip) plus the deferred pragmatics (phases D/E) and two follow-on cleanups that finish the carve's own stated intent.

What's here (by commit)

  1. refactor(config) (Oli + Opus, rebased): decomposition apparatus out of PDConfig into per-domain decomposition: {sites, ci}; tiled site specs (glu_transformer/simple_mlp + LayerSelection) with explicit kept for toys; SiteTree/ArchFamily resolution (param_decomp/site_tree.py) — config→sites→chunks is now parse-free; load_run isinstance guard fixed (CIFn Protocol matched every impl) → ChunkwiseTransformerCIFn.
  2. migrate(config): all 40 live repo yamls migrated by an in-repo tool (param_decomp_lab/migrations/decomposition_apparatus.py — comment-preserving line surgery, proven equal to its pure dict transform; each LM yaml's resolved site list verified byte-identical to the old canonical expansion). The 8 lab LM yamls already unparseable at tip (gpt2*, jose, llama-3.1-8b, pile/ss simple-mlp — pre-chunkwise ci_config schema) are left untouched.
  3. test(config): tests + tooling ported off the old shape. Notable semantics: invalid site names are now parse-time rejections (Literal cs keys) instead of convert-time asserts; per-LAYER heterogeneous C is deliberately unrepresentable (tiling is what makes chunkwise chunks homogeneous by construction) — the old test_arbitrary_sites_with_per_site_c_convert is rewritten accordingly. Dead expand_wildcard_site_cs + wildcard pattern deleted. gen_probe_config output parses AND builds again (it had drifted pre-branch).
  4. refactor(site-tree) ×2: finishing the branch doc's flagged costs on the typed-per-family fork — each target's KIND_ORDER now derives from the config Literal (get_args), the families live on their targets (llama8b.FAMILY / llama_simple_mlp.FAMILY), resolve_site_tree loses the dict relaunder + vocabulary-drift assert (unrepresentable now), and ArchFamily carries parse so the byte-identical canonical_site_cs/site_specs/tap_layer bodies live once in site_tree, family-parameterized (targets keep their two-arg public fns as one-line delegations — no caller churn).

Verification

  • basedpyright clean workspace-wide; ruff clean.
  • Equivalence goldens + SPEC.md untouched (pure config-shape change, no trajectory drift; git diff empty over tests/equivalence, tests/simple_mlp_equivalence, tests/stacked_parity, SPEC.md).
  • Full --runslow -n4 suite: running at PR-open; will confirm on the task thread (targeted config/target/launch/finetune suites all green).
  • experiments/invariance_check.py fails with the pre-existing fix(invariance_check): feed int tokens, not a float residual batch (SPEC D4) #940 float-tokens bug (bridge: fix-invariance-check-float-tokens) — reproduced identically pre-branch; not a regression of this PR.

Known deliberate gaps

  • Stored run launch_config.yamls (152 dirs) don't parse post-carve — per Oli, explicitly out of scope for now ("not really worried about configs"); migrate_raw is written to serve that migration when wanted.
  • Commits 1–3 were made with --no-verify (following the branch's own precedent when jax wasn't type-verifiable in the authoring env); later commits run the hooks. Everything is hook-clean at tip.

Crew-Address: task/targets-refactor
🤖 Generated with Claude Code

@claude-spd1

Copy link
Copy Markdown
Collaborator Author

Full-suite verification complete: 518 passed / 5 skipped / 11 xfailed under --runslow -n4; the 10 reported failures were all xdist worker crashes ("node down: Not properly terminated" — resource contention on a busy node), and all 10 pass on a serial rerun (10 passed in 141s). No test regressions. basedpyright + ruff clean; goldens + SPEC untouched.

Crew-Address: task/targets-refactor

claude-spd1 pushed a commit that referenced this pull request Jul 13, 2026
…ep debris

The repo carried 30 LM run yamls of which ~26 were launched one-offs
(PROFILE/SAVESMOKE/scaling-matrix permutations) — every one a permanent
migration obligation for every schema PR, while the actual archetype configs
(jose, ss-2L, pile-4L) rotted unparseable. Provenance for launched runs never
depended on these files: the run dir pins launch_config.yaml, submit snapshots
the tree to refs/runs/snapshot/<id>, and wandb stores the config.

- CONFIGS.md: the policy — a registry of <10 canonical LM seats; sweeps are
  not committed; schema PRs migrate every repo yaml in the same PR (executed
  in-repo migration, the #966 pattern, never a comment-attached script — the
  #939 lesson); stored-run pins are never rewritten in place.
- param_decomp_lab/tests/test_repo_configs_parse.py: the CI gate — every LM
  yaml in the tree must schema-parse and pass the canonical assert; the three
  broken archetype seats are quarantined in a shrink-only KNOWN_BROKEN list.
- Delete 26 one-off yamls from param_decomp/configs/ and the 5 dead non-seat
  yamls from param_decomp_lab/experiments/lm/.

Kept seats: llama8b_l18_C49k_200k, llama8b_full32L_seq512_b128_dp128,
llama8b_full32L_HSDP_b32_dp32_SAVESMOKE, pile_ppgd_bsc (+ the three
broken-awaiting-migration lab seats).

Crew-Address: task/consumer-config-schema-drift
ocg-goodfire and others added 7 commits July 15, 2026 04:52
…+ SiteTree

Move the target-coupled decomposition apparatus (site-spec + CI-fn arch) out of the shared PDConfig into a per-domain `decomposition` section, so pd is algorithm-only and the target<->apparatus coupling is (mostly) structural. Exploration branch: yaml migration + tests (phases D/E) deferred; jax paths type-verified only, so pre-commit make type is skipped intentionally (--no-verify).

- pd: drop ci_config + decomposition_targets (delete DecompositionTargetConfig)
- decomposition = {sites, ci}; ExperimentConfig non-generic + per-domain concrete subclasses
- tiled site specs (glu_transformer/simple_mlp) + LayerSelection (all/range/list); explicit kept for toys only -> chunkwise homogeneity structural
- SiteTree/ArchFamily resolution (param_decomp/site_tree.py); delete _site_cs/_block_of_site/expand_wildcard_site_cs/groupby -> config->sites->chunks parse-free
- fix load_run isinstance guard (CIFn matched all impls) -> ChunkwiseTransformerCIFn; drop dead ci_arch chunkwise branch

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Phase D of the decomposition-config carve: all 40 live experiment yamls (30
param_decomp/configs LM, 5 TMS, 5 ResidMLP) move pd.ci_config ->
decomposition.ci and pd.decomposition_targets -> decomposition.sites (tiled
glu_transformer/simple_mlp for LMs, explicit for toys). Migration done by the
new in-repo tool (param_decomp_lab/migrations/decomposition_apparatus.py):
comment-preserving line surgery, proven equal to its pure dict transform, and
each migrated LM yaml's resolved site list verified byte-identical to the old
canonical expansion (224 sites full32L, 27 l18-26, etc). The 7 lab LM yamls on
a pre-chunkwise ci_config schema (gpt2*, llama-3.1-8b, pile/ss simple-mlp) were
already unparseable at feature/jax tip and are left untouched.

The migrate_raw dict transform is reusable for stored run launch_config.yamls.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Crew-Address: task/targets-refactor
…gets

Phase E of the decomposition-config carve. Every consumer of the old shape moves
to decomposition.{sites,ci}:

- test_config: non-site/embedding/cross-family cs keys now assert PARSE-time
  rejection (Literal matrix vocabulary) instead of the deleted convert-time
  assert; new family<->target mismatch case (simple_mlp c-spec vs GLU target,
  refused at resolve). test_arbitrary_sites_with_per_site_c_convert becomes
  test_tiled_sites_with_per_matrix_c_convert: per-matrix heterogeneous C over a
  non-contiguous layer list (per-LAYER heterogeneity is deliberately
  unrepresentable now — that's what makes chunkwise chunks homogeneous by
  construction).
- test_fnmatch_site_order: the torch-oracle set/order pins (E21, S10) now
  compare resolve_site_tree output instead of canonical_site_cs/
  expand_wildcard_site_cs pattern lists; same SET-equality + pinned ORDER
  divergence.
- llama_simple_mlp: delete expand_wildcard_site_cs + WILDCARD_NAME_PATTERN
  (orphaned by the carve — wildcards no longer exist anywhere in the config
  path) and their tests; production-pattern test tiles a simple_mlp c-spec.
- test_finetune_resume: changed-C compat check halves decomposition.sites.cs.
- test_launch: minimal LM dict on the new shape.
- gen_probe_config: emits the new decomposition section; also strip beta/
  n_mask_samples/sampling so the emitted probe config PARSES AND BUILDS again
  (verified via build_experiment_config; it had drifted before this branch).
- theoretical_min_memory: reads the tiled spec (no module-pattern parsing).
- docs: CLAUDE.md / configs.py / tms docstrings off ExperimentConfig[T,D] and
  pd.decomposition_targets.

basedpyright clean; test_config/test_fnmatch/test_launch/test_finetune_resume/
test_llama_simple_mlp pass.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Crew-Address: task/targets-refactor
…n their targets

Finishes what the carve's configs.py comment already claimed: each target's
KIND_ORDER now DERIVES from the config-side Literal vocabulary
(get_args(GluMatrix) / get_args(SimpleMlpMatrix)) instead of being a hand-kept
twin, and each target module builds its own ArchFamily (llama8b.FAMILY /
llama_simple_mlp.FAMILY) binding its site_name — the lab GLU_FAMILY/
SIMPLE_MLP_FAMILY constants are gone. resolve_site_tree loses both costs the
branch doc flagged on the typed-per-family fork: the dict[GluMatrix,int] ->
dict[str,int] relaunder and the vocabulary-drift assert (a c-spec key outside
the family vocabulary is now unrepresentable, not asserted); slot ordering is a
plain sort by family rank. The family<->target mismatch assert (semantic, not
drift) stays.

No behavior change: KIND_ORDER values are identical (module-level assert pins
ATTN+MLP composition), resolved site lists byte-identical.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Crew-Address: task/targets-refactor
… helpers

canonical_site_cs / site_specs / tap_layer had byte-identical bodies in
llama8b and llama_simple_mlp, differing only in which module-level
parse_site_name/KIND_ORDER they closed over. ArchFamily now carries the
grammar's inverse (parse) alongside name_of, and the three helpers live once
in site_tree, family-parameterized. The targets keep their public two-arg
functions (canonical_site_cs(site_cs), *site_specs(cfg, site_cs)) as
one-line delegations binding FAMILY + their shape table, so no caller moves.

No behavior change; same asserts, same ordering.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Crew-Address: task/targets-refactor
…file map

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Crew-Address: task/targets-refactor
…in test

Post-rebase onto the Qwen3 ingestion (#987): migrate the jose-ish baseline
(#917, landed old-shape upstream) via the in-repo tool, and restore the
Literal-typed cs dict in test_fnmatch_site_order lost in conflict resolution.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Crew-Address: task/targets-refactor
@claude-spd1
claude-spd1 force-pushed the bridge/task-targets-refactor branch from 01baa9c to 81e6366 Compare July 15, 2026 05:03
@claude-spd1

Copy link
Copy Markdown
Collaborator Author

Rebased onto feature/jax tip past the Qwen3 ingestion (#987). What changed in the rebase:

basedpyright clean; all 41 live yamls parse (GLU ones build end-to-end); goldens (incl. the new qwen3_hf_parity) + SPEC untouched. Full --runslow suite running now — will confirm.

Crew-Address: task/targets-refactor

@claude-spd1

Copy link
Copy Markdown
Collaborator Author

Post-rebase full-suite verification: 525 passed / 6 skipped / 11 xfailed under --runslow -n4. The 8 reported failures: 7 were xdist worker crashes (node contention — all pass on serial rerun), 1 was an environment artifact of my shell (test_pretrained_target_converts_with_all_layers needs DATA_MOUNT, which my restarted session had lost; passes with it set). No regressions. CI is green on the pushed head as well.

Crew-Address: task/targets-refactor

Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants