feat(targets): ingest Qwen/Qwen3-8B-Base as a decomposition target by claude-spd1 · Pull Request #987 · goodfire-ai/param-decomp

claude-spd1 · 2026-07-14T19:44:30Z

Description

Adds Qwen/Qwen3-8B-Base as a decomposition target on feature/jax — JAX fork, direct unit tests against the HuggingFace implementation, and the full DecomposedModel interface, so it can be targeted for a training run.

Qwen3-8B is structurally Llama-3.1-8B plus exactly one delta: per-head RMSNorm on q/k after projection, before RoPE (HF Qwen3Attention.q_norm/k_norm). Same GQA (32h/8kv, head_dim 128), same SiLU GLU MLP, untied embeddings, no biases, plain RoPE (theta 1e6).

Structure (revised per review — model families in their own files, shared machinery honestly labelled, no switches over families):

param_decomp/targets/glu_transformer.py (+ glu_transformer_sharding.py) — the SHARED HF GLU-transformer machinery: site grammar, FrozenAttn/GLULayer/GLUDecomposedModel (renamed from Llama*), the scan/masked-forward engine, and HF safetensors loading parameterized by the family's attn loader + inv_freq. FrozenAttn exposes a _prep_qk pre-RoPE hook (identity) and a target-owned pattern recipe; GLUDecomposedModel exposes attn_pattern(q_site, q_flat, k_flat).
param_decomp/targets/llama8b.py — the Llama family file: llama31_8b_config() (vendored LlamaConfig, llama3 rope scaling — vendored_jax/llama.py untouched) + load_decomposed_llama_from_hf.
param_decomp/targets/qwen3_8b.py — the Qwen3 family file: Qwen3FrozenAttn(FrozenAttn) with REQUIRED q_norm/k_norm/eps fields overriding _prep_qk (QK-norm keyed by attn class, never a config flag — invalid states unrepresentable), qwen3_8b_config() on a family-neutral GLUConfig, load_decomposed_qwen3_from_hf.
attn_patterns_eval.py INVERTED: the pattern recipe is target-owned behind a @runtime_checkable AttnPatternModel protocol; the eval imports no concrete target and the match-on-class dispatch is gone (this also fixes a latent correctness issue: Qwen3's norm weights are per-layer, which a single config-derived closure could not express).
Lab: HF_MODEL_FAMILIES registry in experiments/lm/config.py (model_name → arch config / loader / model_type) — the ONLY place a model name selects a family; build_target / run_metadata consume it. "Llama"/"Qwen3" path schemas registered (the "Llama" entry was missing entirely).
No run yaml committed (per the canonical-seat config policy, PR config: canonical-seat policy, CI parse gate, sweep-debris cleanup #983) — the launch config (252 sites mirroring the llama flagship, 63 sites/chunk → 4 chunks) is attached to task:ingest-qwen-3-8b and validated against the schema in this PR's tests.

Remaining prerequisite for an actual launch (not in this PR): a Qwen3-tokenized prestaged dataset — one command on the training cluster:
python -m param_decomp_lab.experiments.lm.prestage_tokenized --out_dir $PARAM_DECOMP_OUT_DIR/datasets/fineweb_qwen3_tok_512 --tokenizer_name Qwen/Qwen3-8B-Base --seq_len 512

Motivation and Context

task:ingest-qwen-3-8b — Oli wants Qwen3-8B-Base targetable for a training run.

How Has This Been Tested?

NEW direct JAX-vs-HF parity suite (param_decomp/tests/qwen3_hf_parity/) — the first direct HF check in the repo (Llama's chain was JAX↔torch-vendored↔HF):
- tiny-random Qwen3ForCausalLM golden (fp32, eager attention, committed npz + torch-env generator): JAX logits match at rtol 2e-4 / atol 1e-5 ✅
- slow real-weights test: load_decomposed_qwen3_from_hf bf16 forward vs committed HF bf16 final-position logits — KL < 5e-3 + tie-aware argmax ✅ (the golden had a genuine 3-way bf16 top-logit tie on one prompt)
param_decomp/tests/test_qwen3_8b.py: tiny-target contract tests (clean path, mask=1 identity through the QK-norm, ablation, one full train step), QK-norm load-bearing test, and a per-layer QK-norm attn-pattern test
Full make test-all (slow + multidevice) green post-refactor: 533 passed / 7 skipped / 11 xfailed + 5 multidevice (SLURM job on and-gmi)
make check clean (basedpyright 0 errors, ruff clean)
Goldens generated with torch 2.13.0+cpu / transformers 5.13.1 in a throwaway venv (repo stays torch-free)

Does this PR introduce a breaking change?

Renames, no behavior change: LlamaDecomposedModel/LlamaLayer → GLUDecomposedModel/GLULayer, llama_site_specs → glu_site_specs, llama8b_sharding → glu_transformer_sharding, load_decomposed_lm_from_hf → per-family loaders. Llama numerics bit-identical; all importers updated in-repo.

Crew-Address: task/ingest-qwen-3-8b

🤖 Generated with Claude Code

Qwen3-8B-Base is structurally Llama-3.1-8B plus one delta — per-head RMSNorm on q/k after projection, before RoPE (HF Qwen3Attention q_norm/k_norm) — so it rides the existing llama8b target rather than a fork: - `LlamaConfig.qk_norm` (vendored_jax) + optional `FrozenAttn.q_norm/k_norm` (+ static `eps`), applied in `core` after the head reshape, exactly as HF does. Llama paths are untouched (fields default None; SPEC unchanged — q/k sites still decompose before RoPE/SDPA, the masked site output now feeds the norm first). - `qwen3_8b_config()` (36L, d 4096, 32h/8kv, inter 12288, vocab 151936, theta 1e6, eps 1e-6, plain RoPE, untied embed) + `hf_model_config()` model-name allowlist; HF loader reads the q/k_norm keys. - lab dispatch: `_resolve_target` accepts transformers.Qwen3ForCausalLM, d_resid/build_target/run_metadata resolve arch by model name; "Llama"/ "Qwen3" path schemas registered (raw-HF `layers.{i}.self_attn.*` grammar). - attn-patterns eval recipe takes the q-site name and applies the layer's QK-norm (per-layer weights — one shared closure was Llama-only). - tests: qwen-tiny parametrizations of the target suite + a QK-norm load-bearing test; NEW direct JAX-vs-HF parity suite (`tests/qwen3_hf_parity/`): tiny-random Qwen3ForCausalLM golden at fp32 tolerance (passes at rtol 2e-4) and a slow real-weights check (bf16; final-logits KL < 5e-3 + tie-aware argmax; passes against the cluster snapshot). Goldens regenerate via the torch-env `gen_hf_fixtures.py` (torch 2.13.0+cpu, transformers 5.13.1). - launchable config `qwen3_8b_full36L_HSDP_b32_dp32.yaml` mirroring the llama flagship (252 sites, 63 sites/chunk -> 4 chunks). Needs a Qwen3-tokenized prestaged dataset (fineweb_qwen3_tok_512) before launch. Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

ocg-goodfire · 2026-07-14T22:25:52Z

    )


+def qwen3_8b_config() -> LlamaConfig:


what the heck? qwen 3 8b is llama shaped? let's not have sibling imports, if there's genuinely shared utils, let's factor them into an accurately labelled file

ocg-goodfire · 2026-07-14T22:28:09Z

+        q = q_flat.reshape(b, t, self.n_head, self.head_dim)
+        k = k_flat.reshape(b, t, self.n_kv_head, self.head_dim)
+        if self.q_norm is not None:
+            # Qwen3 QK-norm: per-head RMSNorm over head_dim, before RoPE (HF
+            # `Qwen3Attention`). A masked q/k site output feeds this norm, then RoPE/SDPA.
+            assert self.k_norm is not None
+            q = rms_norm(q, self.q_norm, self.eps)
+            k = rms_norm(k, self.k_norm, self.eps)
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)


feels dangerous that we're implementing qwen as switch behaviour inside llama - in general I'd far prefer new code cleanly separated into model family files, over any tricky reuse

ocg-goodfire · 2026-07-14T22:30:15Z

-    def attn_pattern(q_flat: Array, k_flat: Array) -> Array:
+    def attn_pattern(q_site: str, q_flat: Array, k_flat: Array) -> Array:
        b, t, _ = q_flat.shape
        assert q_flat.shape[-1] == n_head * head_dim, q_flat.shape
        assert k_flat.shape[-1] == n_kv_head * head_dim, k_flat.shape
-        q = q_flat.reshape(b, t, n_head, head_dim).transpose(0, 2, 1, 3)
-        k = k_flat.reshape(b, t, n_kv_head, head_dim).transpose(0, 2, 1, 3)
+        q = q_flat.reshape(b, t, n_head, head_dim)
+        k = k_flat.reshape(b, t, n_kv_head, head_dim)
+        if qk_norm_for is not None:
+            q_norm, k_norm = qk_norm_for(q_site)
+            q = rms_norm(q, q_norm, eps)
+            k = rms_norm(k, k_norm, eps)
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
        cos, sin = rope_cos_sin(inv_freq, t, q_flat.dtype)
        q, k = apply_rope(q, k, cos, sin)
        k = repeat_kv(k, n_rep)


this makes me uneasy, as we start having inline switches over model families lets actively reconsider our code structure here. feels like this might already deserve an inversion

ocg-goodfire · 2026-07-14T22:30:37Z

above I said:

this makes me uneasy, as we start having inline switches over model families lets actively reconsider our code structure here. feels like this might already deserve an inversion

same for this whole file honestly

ocg-goodfire · 2026-07-14T22:32:00Z

 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/llama8b.py` + `targets/llama8b_sharding.py` the first target. `targets/llama8b.py` implements BOTH vendored HF GLU-transformer targets — Llama-3.1-8B and Qwen3-8B-Base (`qwen3_8b_config`; the one structural delta is `LlamaConfig.qk_norm`, per-head RMSNorm on q/k before RoPE, optional `FrozenAttn.q_norm`/`k_norm` fields) — dispatched by HF model name via `hf_model_config`. 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


yea this feels like a good sign to refactor

ocg-goodfire · 2026-07-14T22:34:35Z

+    # Qwen3-style per-head RMSNorm on q/k after projection, before RoPE
+    # (HF `Qwen3Attention.q_norm`/`k_norm`). False => the plain Llama attention.
+    qk_norm: bool = False


again, here I'm down to fork at a higher level than to produce objects with many identities where invalid representations become possible due to cartesian products

claude-spd1 · 2026-07-14T22:38:02Z

Agreed on all counts — restructuring now. The plan:

vendored_jax/llama.py reverted — no qk_norm flag on LlamaConfig, no cartesian-product configs.
New param_decomp/targets/glu_transformer.py — the genuinely shared machinery, honestly labelled (matches the glu_transformer family vocabulary from the config-refactor branch): the site grammar, FrozenAttn/GLULayer/GLUDecomposedModel (the scan + masked-forward engine, formerly LlamaDecomposedModel), the HF safetensors reader, and a shared loader parameterized by the family's attn loader. Also a family-neutral GLUConfig for archs with nothing extra.
targets/llama8b.py shrinks to the Llama family file: llama31_8b_config() (vendored LlamaConfig with its llama3 rope scaling) + the Llama attn loader.
New targets/qwen3_8b.py — the Qwen3 family file: Qwen3FrozenAttn(FrozenAttn) with REQUIRED q_norm/k_norm/eps fields (invalid states unrepresentable — QK-norm presence is keyed by attn class, not a flag) overriding a _prep_qk hook, qwen3_8b_config(), and its own HF loader.
attn_patterns_eval.py inverted: the pattern recipe becomes target-owned — attn_pattern(q_site, q_flat, k_flat) on the decomposed model (delegating to the layer's own attn module, which knows its norms), consumed via a @runtime_checkable Protocol like ComponentActivationModel. The eval file loses its imports of concrete targets and the match-on-class dispatch entirely.
Lab keeps the model-name → family dispatch (config/loader/model_type registry in experiments/lm/config.py), so core never holds a cross-family switch.
llama8b_sharding.py → glu_transformer_sharding.py (it was already family-generic).

Same tests (incl. the direct HF-parity suite) re-pointed; qwen tiny-target tests move to their own test_qwen3_8b.py.

Crew-Address: task/ingest-qwen-3-8b

…er target (review) Per PR #987 review: no switch-behavior inside llama8b, no cartesian-product configs, no inline model-family dispatch in evals. - targets/glu_transformer.py (+ _sharding): the SHARED HF GLU-transformer machinery — site grammar, FrozenAttn/GLULayer/GLUDecomposedModel (renamed from Llama*), scan/masked-forward engine, HF loading parameterized by the family's attn loader + inv_freq. FrozenAttn gains a `_prep_qk` pre-RoPE hook (identity) and a target-owned `pattern` recipe; GLUDecomposedModel gains `attn_pattern(q_site, q_flat, k_flat)`. - targets/llama8b.py: the Llama family file — llama31_8b_config (vendored LlamaConfig, llama3 rope) + load_decomposed_llama_from_hf. - targets/qwen3_8b.py: the Qwen3 family file — Qwen3FrozenAttn with REQUIRED q_norm/k_norm/eps fields overriding _prep_qk (QK-norm keyed by attn class, never a flag; invalid states unrepresentable), qwen3_8b_config (GLUConfig), load_decomposed_qwen3_from_hf. vendored_jax/llama.py reverted untouched. - attn_patterns_eval inverted: the recipe is target-owned behind a runtime_checkable AttnPatternModel protocol; the eval no longer imports any concrete target and the match-on-class dispatch is gone. Step factories drop the pattern_fn param. SimpleMLPDecomposedModel implements the method. - lab: HF_MODEL_FAMILIES registry (model_name -> arch_config/load/model_type) in experiments/lm/config.py — the ONLY place a model name selects a family; build_target/run_metadata consume it. - qwen tiny-target tests move to test_qwen3_8b.py (incl. per-layer QK-norm pattern test); parity suite re-pointed at the family loader. Full local run of affected suites green; make check clean. Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

claude-spd1 · 2026-07-14T23:09:31Z

Restructure pushed per the review: model families in their own files (targets/llama8b.py, targets/qwen3_8b.py), shared machinery factored into targets/glu_transformer.py (+_sharding), vendored_jax/llama.py reverted untouched, QK-norm keyed by attn class with REQUIRED fields (no flag), and attn_patterns_eval inverted behind a target-owned attn_pattern protocol — no model-family switches left outside the lab's HF_MODEL_FAMILIES registry. Full make test-all green post-refactor (533 passed + multidevice); make check clean; CI passing. Ready for another look.

Crew-Address: task/ingest-qwen-3-8b

The launch config is a one-off run recipe, not a canonical seat — it rides the task thread (task:ingest-qwen-3-8b), and a launched run pins its own copy in the run dir. Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

ocg-goodfire · 2026-07-14T23:19:02Z

+from param_decomp.sharding import assert_divisible
+from vendored_jax.llama import apply_rope, causal_sdpa, repeat_kv, rms_norm, rope_cos_sin
+
+DT = jnp.bfloat16


the fact that this is global to a shared "glu_transformer" module seems wrong to me. seems overly opinionated for an abstract module and sure maybe it works for now, but seems like that's only specifically because we happen to currently only need qwen and llama

Agreed — dropped the module-global: HFWeights now takes the dtype, load_decomposed_glu_from_hf threads it, and each FAMILY loader passes bf16 (the thing that actually holds that opinion, matching its supported_weights_dtypes). The shared module no longer carries a dtype at all.

…hared-module global Per review: drop `DT = jnp.bfloat16` from glu_transformer — HFWeights takes the dtype, load_decomposed_glu_from_hf threads it, and each family loader passes bf16 (matching its TargetConfig.supported_weights_dtypes). Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…eutral docstring - HFModelFamily gains model_class; _resolve_target derives the accepted set from the registry and cross-checks class <-> model_name (was a parallel hardcoded tuple — a second place enumerating families). - targets/target_aliases.py deleted: zero importers, and its docstring claimed a dispatch that no longer exists. - GLUDecomposedModel docstring stops enumerating concrete families. Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ferences) Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… forward per step GLUDecomposedModel.read_activations was a Python loop over all blocks — the largest unrolled region left in the step HLO, present in the train step and all seven jitted eval/metric steps — and a duplicate of the scanned clean_output forward on the same batch two lines earlier (XLA does not CSE scan vs unroll). (Authored against llama8b.py pre-#987; the machinery now lives in glu_transformer.py, so the fix covers llama8b AND qwen3.) Now clean_output / read_activations / the new clean_output_and_activations all route through one frozen-path lax.scan (S3/S4 semantics unchanged): taps ride out as scan ys, one stacked [tapped_depth,b,t,d] array per block intermediate (resid / post-LN1 / attn-out / post-LN2 / MLP-inner), indexed into the flat tap dict after the scan — the _run_masked_forward per-kind-ys pattern, no per-layer lax.cond. The tapped prefix alone emits ys, so the old early-break compute property is preserved statically. Family behavior (Qwen3 QK-norm) rides in attn.core exactly as the unrolled version had it. clean_output_and_activations (DecomposedModel protocol; trivial two-call impl on the unscanned targets) serves the two steps that need the clean logits AND the taps on the same batch — train and fast eval — from ONE forward; harvest's load_run forward folds three frozen passes into it likewise. Numerics (pinned pre-refactor CPU fixtures, 557 arrays): clean_output logits BIT-identical; taps vs the old unrolled forward differ within fp32 reassociation (<=1.7e-5 abs — the scan-vs-unroll class clean_output's docstring already documents; same order as SPEC D4's accepted 2.2e-6), propagating to <=3.9e-6 over a 3-step trajectory in both remat modes. The fused accessor is exact vs the separate methods (pinned by the new test); taps are now bit-identical to the residual the recon target actually threads. Crew-Address: task/scan-read-activations-taps

…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

ocg-goodfire suggested changes Jul 14, 2026

View reviewed changes

PD User (shared) and others added 2 commits July 14, 2026 22:56

docs: point README/CLAUDE.md target tables at the family split

2e82e3c

Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

ocg-goodfire reviewed Jul 14, 2026

View reviewed changes

PD User (shared) and others added 3 commits July 14, 2026 23:50

style: drop needless quoted annotations (quotes only for true self-re…

9404685

…ferences) Crew-Address: task/ingest-qwen-3-8b Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

ocg-goodfire merged commit a2c297b into feature/jax Jul 15, 2026
1 check passed

claude-spd1 mentioned this pull request Jul 15, 2026

perf(glu): emit residual taps from the clean-forward scan; one frozen forward per step #972

Open

claude-spd1 mentioned this pull request Jul 15, 2026

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

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

feat(targets): ingest Qwen/Qwen3-8B-Base as a decomposition target#987
ocg-goodfire merged 7 commits into
feature/jaxfrom
bridge/task-ingest-qwen-3-8b

claude-spd1 commented Jul 14, 2026 •

edited

Loading

Uh oh!

ocg-goodfire Jul 14, 2026

Uh oh!

ocg-goodfire Jul 14, 2026

Uh oh!

ocg-goodfire Jul 14, 2026

Uh oh!

ocg-goodfire Jul 14, 2026

Uh oh!

ocg-goodfire Jul 14, 2026

Uh oh!

ocg-goodfire Jul 14, 2026

Uh oh!

claude-spd1 commented Jul 14, 2026

Uh oh!

claude-spd1 commented Jul 14, 2026

Uh oh!

ocg-goodfire Jul 14, 2026

Uh oh!

claude-spd1 Jul 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

claude-spd1 commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

How Has This Been Tested?

Does this PR introduce a breaking change?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude-spd1 commented Jul 14, 2026

Uh oh!

claude-spd1 commented Jul 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

claude-spd1 commented Jul 14, 2026 •

edited

Loading