feat(targets): ingest Qwen/Qwen3-8B-Base as a decomposition target

CLAUDE.md

@@ -55,7 +55,7 @@ lab is composition / IO / CLI / experiment assembly**:
 - **`param-decomp`** (root: `param_decomp/` + sibling `pretrain/` + sibling
   `vendored_jax/`) — the core: the generic JAX single-pool VPD trainer ENGINE
   (`param_decomp/`: run.py = `run_decomposition_training`, lm.py, train.py, ci_fn.py,
-  targets/llama8b.py, …), the torch-free pydantic config SCHEMA it now carries directly
+  targets/glu_transformer.py + its llama8b/qwen3_8b family files, …), the torch-free pydantic config SCHEMA it now carries directly
   (`base_config.py` = `BaseConfig`, `schedule.py`, `configs.py` = `PDConfig` /
   `RuntimeConfig` / `Cadence` / loss + eval-metric configs / routing / ci-fn / wandb
   shaping), the built-run bundle (`built_run.py`: `BuiltRun` / `DataConfig` /
@@ -157,7 +157,7 @@ and returns JAX-native as the #10 torch->jax adapter.
   built-run bundle the engine consumes (`built_run.py`: `BuiltRun` /
   `DataConfig` / `EvalConfig` / … + the `TargetSites` protocol). The engine + numerics
   (`run.py` = `run_decomposition_training`, `lm.py` / `train.py` / `ci_fn.py` /
-  `targets/llama8b.py` / `targets/llama_simple_mlp.py` / `adversary.py` / `recon.py` / `losses.py` /
+  `targets/glu_transformer.py` (+ `targets/{llama8b,qwen3_8b}.py` family files) / `targets/llama_simple_mlp.py` / `adversary.py` / `recon.py` / `losses.py` /
   `checkpoint.py` / `sharding.py` / `eval.py` / `slow_eval.py` / `arithmetic_eval.py` +
   `log.py`) plus `configs/`
   (the self-contained run yamls) and `tests/` (incl. the `tests/equivalence/` frozen

param_decomp/CLAUDE.md

@@ -27,7 +27,7 @@ never silently diverge. Cite IDs (`S14`, `N1`, …) in commit messages and revie
 `lm.py` defines `DecomposedModel` — a `@runtime_checkable Protocol`: ordered `sites` +
 `leading_axes` + the methods `clean_output`, `read_activations`, `masked_output`,
 `masked_site_outputs`, `weight_deltas`, and a `recon_loss_fn` (LM: `kl_per_position`). The
-concrete impl per target is an `eqx.Module` (`LlamaDecomposedModel`,
+concrete impl per target is an `eqx.Module` (`GLUDecomposedModel`,
 `SimpleMLPDecomposedModel`, `TMSDecomposedModel`, `ResidMLPDecomposedModel`) carrying its
 FROZEN target weights as ARRAY FIELDS; the TRAINABLE V/U (`vu: DecompVU`) stays an explicit
 METHOD ARG (separate lifecycle — own optimizer + checkpoint, C-sharded while the frozen
@@ -58,12 +58,12 @@ single `make_plan` constructor, built from the shared configs by `recon.build_lo
 see LOSS_PARITY_DESIGN.md),
 consuming `losses.py` (pure loss terms + schedules) and `adversary.py` (persistent
 vs fresh source machinery — semantically distinct adversaries sharing only
-`source_masks`); `ci_fn.py` the shared CI transformer; `targets/llama8b.py` + `targets/llama8b_sharding.py` the first target. There is ONE
+`source_masks`); `ci_fn.py` the shared CI transformer; `targets/glu_transformer.py` + `targets/glu_transformer_sharding.py` the SHARED HF GLU-transformer target machinery (site grammar, `FrozenAttn`/`GLULayer`/`GLUDecomposedModel`, the scan/masked-forward engine, HF loading), with the model FAMILIES in their own files: `targets/llama8b.py` (vendored `LlamaConfig`, llama3 rope) and `targets/qwen3_8b.py` (`Qwen3FrozenAttn` — REQUIRED `q_norm`/`k_norm` fields applied in the `_prep_qk` pre-RoPE hook; Qwen3's one structural delta). Nothing in the shared file switches on a family; the model-name → family registry is LAB-side (`experiments/lm/config.py::HF_MODEL_FAMILIES`). Qwen3 JAX↔HF parity is pinned DIRECTLY by `tests/qwen3_hf_parity/` (a tiny-random `Qwen3ForCausalLM` golden at fp32 tolerance + a slow real-weights logits check; goldens regenerate via its torch-env `gen_hf_fixtures.py`). There is ONE
 recon semantics: masks thread through the full token-input forward, loss is KL on final logits
 (SPEC §2.3–2.5). Site-local recon is a conceptual no-no, not a "simplification".
 `targets/llama_simple_mlp.py` is the second target (the pile-pretrained `LlamaSimpleMLP`,
 t-9d2b8f02; sites `h.{i}.attn.{q,k,v,o}_proj` / `h.{i}.mlp.{c_fc,down_proj}`) —
-config dispatch is `TargetConfig` (llama8b) vs `LlamaSimpleMLPTargetConfig`, both LAB-side
+config dispatch is `TargetConfig` (the HF GLU families) vs `LlamaSimpleMLPTargetConfig`, both LAB-side
 (`param_decomp_lab/experiments/lm/config.py`, which reads the canonical schema DIRECTLY —
 `build_experiment_config`/`load_config` — routing `kind: pretrained` specs + `h.*`
 wildcards), target build in the LM composition root
@@ -98,7 +98,7 @@ every rank builds the identical grid, no rank-0 write or barrier), so configs st
 cluster-portable. The ONE fused `make_arithmetic_grid_step` slices, at the answer position with
 the BATCH axis KEPT as the grid, each component's lower-leaky CI (from the CI fn) and its
 pre-mask activation `x@V` (from the decomposed forward under all-ones masks — the
-`masked_component_activations` seam, llama8b-only, narrowed via the `ComponentActivationModel`
+`masked_component_activations` seam, GLU-target-only, narrowed via the `ComponentActivationModel`
 Protocol). The device→host pull is TWO-PHASE (`compute_arithmetic_selection`), sized to what
 the figures need — never the full `(n_prompts, C)` grids (~GBs/site at production C): the
 step's replicated per-component max CI (over REAL rows only — the sharding-pad tail is
@@ -155,7 +155,7 @@ glued-ΣC head mid-site). **ZeRO-1 ÷N**: the trainable V/U + CI-fn fp32 masters
 m/v shard ÷N over the FULL mesh (`("replicate","fsdp")` on V's d_in / U's d_out / the CI fn's
 d_model) — the dominant optimizer-state memory scales 1/N, not the fixed 1/fsdp. The bf16
 COMPUTE weights are reconstructed to the `fsdp`-sharded (÷fsdp) layout ONCE per step in ENTRY
-(the cross-`replicate` gather, off the hot path — `llama8b._reconstruct_compute_weights` /
+(the cross-`replicate` gather, off the hot path — `glu_transformer._reconstruct_compute_weights` /
 `ci_fn._reconstruct_ci_compute_weights` pin `P(None,"fsdp",...)` BEFORE the per-layer /
 per-chunk scan), landing a SMALL ÷fsdp-resident stack; the scan body then gathers ONE layer's
 `fsdp` shard to full d_in transiently (NVLink, freed each iteration) — NEVER a

param_decomp/README.md

@@ -31,7 +31,8 @@ root with sibling packages `pretrain/` (the in-house target-LM pretrainer) and
 | `run_state.py` | optimizer + initial-`TrainState` construction from an `ExperimentConfig` (orbax restores onto this reference) |
 | `tools/` | `convert_llama_simple_mlp_checkpoint.py` (torch venv) — one-off `.pt` → safetensors conversion of the pile pretrain checkpoint; `migrate_c49k_checkpoint.py` — one-off remap of the frozen C49k clone's orbax `TrainState` (legacy `components.{Vg..Ud}` `(1,*,*)` + flat `sources.<site>`) onto the current layout (site-keyed `components.vu`, `sources.<state_key>.<site>`) so a fine-tune can `restore_latest` it |
 | `sharding.py` | generic GSPMD helpers (`init_distributed`, `dp_mesh`, `replicate`, `shard_batch`) |
-| `targets/llama8b.py` | Llama-3.1-8B target (`LlamaDecomposedModel`, full-model token-input forward, embed internal), arbitrary per-layer matrix sites (`q/k/v/o/gate/up/down`, per-site C; q/k/v decomposed before RoPE/SDPA), per-site `DecompVU`, HF safetensors loader (`build_decomposed_lm` / `load_decomposed_lm_from_hf`) |
+| `targets/glu_transformer.py` | the shared HF GLU-transformer target machinery (`GLUDecomposedModel`, full-model token-input forward, embed internal), arbitrary per-layer matrix sites (`q/k/v/o/gate/up/down`, per-site C; q/k/v decomposed before `_prep_qk`/RoPE/SDPA), per-site `DecompVU`, HF safetensors loading (`build_decomposed_lm` / `load_decomposed_glu_from_hf`) |
+| `targets/llama8b.py` / `targets/qwen3_8b.py` | the model FAMILY files: arch config + `FrozenAttn` variant (Qwen3: `Qwen3FrozenAttn` QK-norm) + HF attn loader |
 | `targets/llama_simple_mlp.py` | `LlamaSimpleMLP` pile-pretrained target (`goodfire/spd/runs/t-9d2b8f02`: 4L, d768, GELU MLP, plain rotate-half RoPE, tied head): sites `h.{i}.attn.{q,k,v,o}_proj` / `h.{i}.mlp.{c_fc,down_proj}` with `h.*` wildcard expansion, pretrain-cache safetensors loader (one-off `.pt` conversion: `tools/convert_llama_simple_mlp_checkpoint.py`), `llama_simple_mlp_decomposed_lm(cfg, sites)`; frozen weights small enough to replicate (`replicate_frozen`), V/U/CI/source placement reuses the generic per-site plan |
 | `run.py` | the generic ENGINE `run_decomposition_training` (pure library, no `main`/YAML): faith warmup, loop, metrics jsonl/wandb, in-loop slow renderer, orbax checkpoints, SIGTERM-save + requeue-resume. The LM composition root that reads YAML + builds the target lives lab-side (`param_decomp_lab/experiments/lm/run.py`) |
 | `data.py` | deterministic batch schedule over the pre-tokenized fineweb parquet shards; O(1) resume addressing, per-process slices |
@@ -41,7 +42,7 @@ root with sibling packages `pretrain/` (the in-house target-LM pretrainer) and
 | `base_config.py` | `BaseConfig` (frozen `extra=forbid` pydantic `BaseModel` + YAML/JSON round-trip), `Probability` |
 | `schedule.py` | `ScheduleConfig` + its two evaluators, host `get_scheduled_value` / traced `scheduled_value_traced` (warmup → constant/linear/cosine decay; every scheduled quantity routes through here) |
 | `configs/` | the single self-contained run yamls (one file per run; no wrapper/schema split) |
-| `targets/llama8b_sharding.py` | the 8B placement plan (frozen replicated; per-site V/U + CI + Adam C-sharded; source replicated; batch sharded) |
+| `targets/glu_transformer_sharding.py` | the 8B placement plan (frozen replicated; per-site V/U + CI + Adam C-sharded; source replicated; batch sharded) |
 | `experiments/llama8b_real.py` | the runnable 8B step + tok/s/GPU bench |
 | `experiments/invariance_check.py` | device-count invariance harness (SPEC D4) |
 | `tests/` | tiny-target unit tests (incl. attention sites + heterogeneous per-site C), checkpoint resume, sharding, `tests/equivalence/` — the fixture-driven torch↔JAX loss-term equivalence harness — `tests/stacked_parity/` — fixtures pinning the pre-site-generality stacked implementation (clean logits bit-identical, train trajectory rel ≤ ~1e-5) — and `tests/simple_mlp_equivalence/` — torch-fixture logits parity for the LlamaSimpleMLP target (tiny random model max abs diff ~2e-7; real t-9d2b8f02 weights ~5e-5 fp32) |

param_decomp/arithmetic_eval.py

@@ -25,7 +25,7 @@
 @runtime_checkable
 class ComponentActivationModel(DecomposedModel, Protocol):
     """A `DecomposedModel` that also exposes per-component activations `x@V`. The arithmetic
-    activation heatmaps need this seam; it is LM-only (currently `LlamaDecomposedModel`), so
+    activation heatmaps need this seam; it is LM-only (currently `GLUDecomposedModel`), so
     the eval narrows to it with an `isinstance` check rather than widening the core
     `DecomposedModel` Protocol every target must satisfy."""