diff --git a/CLAUDE.md b/CLAUDE.md index fae0420fb..8cb091097 100644 --- a/CLAUDE.md +++ b/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 diff --git a/param_decomp/CLAUDE.md b/param_decomp/CLAUDE.md index 0914b2a4d..173c6ff5e 100644 --- a/param_decomp/CLAUDE.md +++ b/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 diff --git a/param_decomp/README.md b/param_decomp/README.md index 4d4a11de3..507f256fb 100644 --- a/param_decomp/README.md +++ b/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.`) onto the current layout (site-keyed `components.vu`, `sources..`) 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) | diff --git a/param_decomp/arithmetic_eval.py b/param_decomp/arithmetic_eval.py index ef3ab0fba..f546e4a56 100644 --- a/param_decomp/arithmetic_eval.py +++ b/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.""" diff --git a/param_decomp/attn_patterns_eval.py b/param_decomp/attn_patterns_eval.py index 146fd7f6d..0f356736a 100644 --- a/param_decomp/attn_patterns_eval.py +++ b/param_decomp/attn_patterns_eval.py @@ -16,21 +16,21 @@ seam — the same all-false-routes trick `hidden_acts_eval` uses for its clean target (one all-false forward) plus one masked forward — so `DecomposedModel` gains no new method. -The attention-pattern reproduction is target-specific (RoPE base, GQA, head reshape) and -lives HERE, not on `DecomposedModel`: `attn_pattern_for` dispatches on the concrete frozen -target, reusing that target's OWN RoPE helper and attention config. A non-attention target -(no `FrozenAttn`/`inv_freq`) raises — the metric only applies to attention-bearing targets. +The attention-pattern reproduction is target-specific (RoPE base, GQA, head reshape, +Qwen3's per-layer QK-norm), so it is TARGET-OWNED: the model exposes +`attn_pattern(q_site, q_flat, k_flat)` (the `AttnPatternModel` protocol below — +`GLUDecomposedModel` / `SimpleMLPDecomposedModel` delegate to their own attention +modules). This file only drives it; nothing here switches on a model family. A target +without the method refuses at step-build time. Masked and clean Q/K run in COMPUTE_DT (bf16, matching the trained model); the pattern softmax and the KL reduction are fp32. """ -import math from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import Any, Protocol, runtime_checkable -import jax import jax.numpy as jnp import numpy as np from jax import random @@ -39,64 +39,22 @@ from param_decomp.components import DecompVU from param_decomp.jit_util import filter_jit from param_decomp.lm import DecomposedModel, all_false_routes -from param_decomp.targets.llama8b import FrozenAttn, LlamaDecomposedModel -from param_decomp.targets.llama_simple_mlp import SimpleMLPDecomposedModel from param_decomp.train import COMPUTE_DT, cast_floating -from vendored_jax.llama import apply_rope, repeat_kv, rope_cos_sin - -AttnPatternFn = Callable[[Float[Array, "B T qd"], Float[Array, "B T kvd"]], Float[Array, "B H T T"]] -"""`(q_flat, k_flat) -> (B, n_heads, T_query, T_key)` post-softmax causal attention map -from a layer's flat Q/K projections — the per-target attention recipe.""" - - -def _attn_pattern_from_config( - n_head: int, n_kv_head: int, head_dim: int, n_rep: int, inv_freq: Array -) -> AttnPatternFn: - """The shared RoPE + GQA + causal-softmax pattern recipe, parameterised by a target's - attention config. Reuses the vendored `rope_cos_sin`/`apply_rope`/`repeat_kv` — never - a reimplemented RoPE. Scores in fp32, scaled by `1/√head_dim`, causal-masked, softmaxed.""" - - def attn_pattern(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) - 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) - scores = jnp.einsum("bhqd,bhkd->bhqk", q.astype(jnp.float32), k.astype(jnp.float32)) - scores = scores / math.sqrt(head_dim) - causal = jnp.triu(jnp.ones((t, t), bool), k=1) - scores = jnp.where(causal, -jnp.inf, scores) - return jax.nn.softmax(scores, axis=-1) - - return attn_pattern - - -def _frozen_attn(target: LlamaDecomposedModel | SimpleMLPDecomposedModel) -> FrozenAttn: - """The first layer's attention — every layer shares the same attn config, so - one `FrozenAttn` fixes `n_head`/`n_kv_head`/`head_dim`/`n_rep` for the recipe.""" - assert target.layers, "attn-patterns metric needs at least one layer" - return target.layers[0].attn - - -def attn_pattern_for(target: Any) -> AttnPatternFn: - """The per-target attention-pattern recipe, dispatched on the concrete frozen target. - - Reuses the target's OWN attention config (`FrozenAttn`) and RoPE frequencies - (`inv_freq`). A non-attention target raises — the metric only applies to - attention-bearing targets (localize-and-assert: fail loudly, never silently).""" - match target: - case LlamaDecomposedModel() | SimpleMLPDecomposedModel(): - attn = _frozen_attn(target) - return _attn_pattern_from_config( - attn.n_head, attn.n_kv_head, attn.head_dim, attn.n_rep, target.inv_freq - ) - case _: - raise AssertionError( - f"attn-patterns metric only applies to attention targets, got {type(target).__name__}" - ) + + +@runtime_checkable +class AttnPatternModel(Protocol): + """The capability this eval needs from a target: the post-softmax `(B, H, T, T)` + causal attention map from one layer's flat Q/K projections, identified by the + layer's `q_proj` site name (the recipe is the target's own — RoPE flavor, GQA, + any pre-RoPE math like Qwen3's per-layer QK-norm).""" + + def attn_pattern( + self, + q_site: str, + q_flat: Float[Array, "B T qd"], + k_flat: Float[Array, "B T kvd"], + ) -> Float[Array, "B H T T"]: ... def _attn_layer_sites(site_names: tuple[str, ...]) -> tuple[tuple[str, str], ...]: @@ -145,7 +103,6 @@ def _pattern_kl(target_pattern: Array, masked_pattern: Array) -> Array: def _clean_patterns( model: DecomposedModel, - pattern_fn: AttnPatternFn, layer_pairs: tuple[tuple[str, str], ...], prepared: Any, tokens: Int[Array, "*leading"], @@ -162,25 +119,29 @@ def _clean_patterns( {s: jnp.zeros(leading, COMPUTE_DT) for s in site_names}, all_false_routes(site_names, leading), site_names, False, ) # fmt: skip - return {q: pattern_fn(clean_outputs[q], clean_outputs[k]) for q, k in layer_pairs} + assert isinstance(model, AttnPatternModel) + return {q: model.attn_pattern(q, clean_outputs[q], clean_outputs[k]) for q, k in layer_pairs} def _masked_patterns_kl( - pattern_fn: AttnPatternFn, + model: DecomposedModel, layer_pairs: tuple[tuple[str, str], ...], masked_outputs: dict[str, Array], target_patterns: dict[str, Array], ) -> dict[str, Array]: + assert isinstance(model, AttnPatternModel) return { - q: _pattern_kl(target_patterns[q], pattern_fn(masked_outputs[q], masked_outputs[k])) + q: _pattern_kl( + target_patterns[q], model.attn_pattern(q, masked_outputs[q], masked_outputs[k]) + ) for q, k in layer_pairs } def _assert_attention_sequence_axes(lm: DecomposedModel) -> None: """Attention patterns are `(B, H, T_query, T_key)` causal maps over a sequence axis; - the metric only applies to a sequence-axis LM target (complements the per-target - `attn_pattern_for` dispatch, which rejects non-attention targets).""" + the metric only applies to a sequence-axis LM target (the `AttnPatternModel` + capability assert in the step factories rejects non-attention targets).""" assert lm.leading_axes == ("sequence",), ( f"attn-patterns eval is LM-only (causal attention over a sequence axis); model " f"has leading_axes={lm.leading_axes}" @@ -189,12 +150,14 @@ def _assert_attention_sequence_axes(lm: DecomposedModel) -> None: def make_ci_attn_patterns_step( lm: DecomposedModel, - pattern_fn: AttnPatternFn, compiler_options: dict[str, bool | int | str] | None = None, ) -> AttnPatternsStep: """Deterministic CI-mask attn-patterns step: `lower_leaky` CI, no delta, one masked forward + one clean (all-false) forward.""" _assert_attention_sequence_axes(lm) + assert isinstance(lm, AttnPatternModel), ( + f"attn-patterns eval needs a target exposing attn_pattern; {type(lm).__name__} does not" + ) site_names = lm.site_names layer_pairs = _attn_layer_sites(site_names) @@ -211,15 +174,13 @@ def step( ci_fn_bf16 = cast_floating(ci_fn, COMPUTE_DT) ci_lower = ci_fn_bf16(taps, remat=False).lower - target_patterns = _clean_patterns( - model, pattern_fn, layer_pairs, prepared, tokens, ci_lower - ) + target_patterns = _clean_patterns(model, layer_pairs, prepared, tokens, ci_lower) leading = tokens.shape zeros_delta = {s: jnp.zeros(leading, COMPUTE_DT) for s in site_names} masked_outputs = model.masked_site_outputs( prepared, tokens, ci_lower, zeros_delta, None, site_names, False ) - sum_kl = _masked_patterns_kl(pattern_fn, layer_pairs, masked_outputs, target_patterns) + sum_kl = _masked_patterns_kl(model, layer_pairs, masked_outputs, target_patterns) n_distributions = {q: int(np.prod(target_patterns[q].shape[:3])) for q, _ in layer_pairs} return sum_kl, n_distributions @@ -228,7 +189,6 @@ def step( def make_stochastic_attn_patterns_step( lm: DecomposedModel, - pattern_fn: AttnPatternFn, n_mask_samples: int, compiler_options: dict[str, bool | int | str] | None = None, ) -> AttnPatternsStep: @@ -236,6 +196,9 @@ def make_stochastic_attn_patterns_step( (with weight deltas), per-draw per-layer pattern KL summed. RNG via per-draw / per-site `fold_in` (the eval-step discipline, mirrors `hidden_acts_eval`).""" _assert_attention_sequence_axes(lm) + assert isinstance(lm, AttnPatternModel), ( + f"attn-patterns eval needs a target exposing attn_pattern; {type(lm).__name__} does not" + ) assert n_mask_samples >= 1, n_mask_samples site_names = lm.site_names layer_pairs = _attn_layer_sites(site_names) @@ -253,9 +216,7 @@ def step( ci_fn_bf16 = cast_floating(ci_fn, COMPUTE_DT) ci_lower = ci_fn_bf16(taps, remat=False).lower - target_patterns = _clean_patterns( - model, pattern_fn, layer_pairs, prepared, tokens, ci_lower - ) + target_patterns = _clean_patterns(model, layer_pairs, prepared, tokens, ci_lower) leading = tokens.shape sum_kl = {q: jnp.zeros((), jnp.float32) for q, _ in layer_pairs} @@ -274,7 +235,7 @@ def step( masked_outputs = model.masked_site_outputs( prepared, tokens, masks, delta_masks, None, site_names, True ) - draw_kl = _masked_patterns_kl(pattern_fn, layer_pairs, masked_outputs, target_patterns) + draw_kl = _masked_patterns_kl(model, layer_pairs, masked_outputs, target_patterns) sum_kl = {q: sum_kl[q] + draw_kl[q] for q, _ in layer_pairs} n_distributions = { diff --git a/param_decomp/ci_fn.py b/param_decomp/ci_fn.py index a2f8c95da..8c9d9dadf 100644 --- a/param_decomp/ci_fn.py +++ b/param_decomp/ci_fn.py @@ -736,7 +736,7 @@ def init_global_mlp_ci_fn( CIFnArch = ChunkwiseTransformerCIArch | MLPCIArch | GlobalMLPCIArch """Every CI-fn architecture. Construction goes through `build_ci_fn`; sharding/placement is -a separate, scale-driven concern (see `llama8b_sharding`), never coupled to arch type.""" +a separate, scale-driven concern (see `glu_transformer_sharding`), never coupled to arch type.""" def build_ci_fn(arch: CIFnArch, sites: tuple[SiteSpec, ...], key: PRNGKeyArray) -> CIFn: diff --git a/param_decomp/experiments/invariance_check.py b/param_decomp/experiments/invariance_check.py index e30bf712f..5276cd58a 100644 --- a/param_decomp/experiments/invariance_check.py +++ b/param_decomp/experiments/invariance_check.py @@ -49,8 +49,8 @@ from param_decomp.recon import build_loss_terms from param_decomp.schedule import ScheduleConfig from param_decomp.sharding import hsdp_mesh, shard_batch -from param_decomp.targets.llama8b import ( - llama_site_specs, +from param_decomp.targets.glu_transformer import ( + glu_site_specs, mlp_family_site_cs, ) from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm @@ -60,7 +60,7 @@ def _run(steps: int, sharded: bool) -> list[dict[str, float]]: cfg = _tiny_cfg() C, seq, gbatch = 8, 16, 8 - sites = llama_site_specs(cfg, mlp_family_site_cs(3, 6, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(3, 6, C)) lm = _tiny_decomposed_lm(cfg, sites, random.PRNGKey(0)) vu = init_decomp_vu(sites, random.PRNGKey(1)) arch = ChunkwiseTransformerCIArch( diff --git a/param_decomp/hf_http.py b/param_decomp/hf_http.py index 72b53717b..0119dd74e 100644 --- a/param_decomp/hf_http.py +++ b/param_decomp/hf_http.py @@ -1,7 +1,7 @@ """Make huggingface_hub's HTTP backend resilient to a cold-cache startup burst. The JAX trainer normally loads the frozen target's weights from a pre-warmed local -HF snapshot (`llama8b.hf_snapshot_dir` asserts the snapshot exists, no network call). +HF snapshot (`glu_transformer.hf_snapshot_dir` asserts the snapshot exists, no network call). But at production topology (8 ranks/node x N nodes) a *cold* cache turns startup into an 8N-rank simultaneous Hub burst, and a single `ReadTimeout` on one rank tears the whole job down before training begins. This mounts a retrying adapter on huggingface_hub's diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index a20bfc436..98a99375d 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -30,7 +30,7 @@ persistent_configs, ) from param_decomp.schedule import ScheduleConfig -from param_decomp.targets.llama8b_sharding import ( +from param_decomp.targets.glu_transformer_sharding import ( init_ci_fn_placed, init_decomp_vu_placed, init_sources_sharded, diff --git a/param_decomp/sharding.py b/param_decomp/sharding.py index bac78a874..1523fc04b 100644 --- a/param_decomp/sharding.py +++ b/param_decomp/sharding.py @@ -13,7 +13,7 @@ the batch axis. No manual NCCL, no pool-coordination code. Placement is expressed as `NamedSharding`: target-specific plans (like -`llama8b_sharding.py`) FSDP-shard params on `fsdp`; `shard_batch` shards the data axis over +`glu_transformer_sharding.py`) FSDP-shard params on `fsdp`; `shard_batch` shards the data axis over the full mesh. """ diff --git a/param_decomp/targets/glu_transformer.py b/param_decomp/targets/glu_transformer.py new file mode 100644 index 000000000..4c92e896c --- /dev/null +++ b/param_decomp/targets/glu_transformer.py @@ -0,0 +1,1145 @@ +"""The SHARED machinery of the vendored HF GLU-transformer decomposition targets: site +grammar, frozen modules, the scan/masked-forward engine, and HF safetensors loading. The +model FAMILIES live in their own files — `targets/llama8b.py`, `targets/qwen3_8b.py` — +each contributing its arch config, its `FrozenAttn` variant (via the `_prep_qk` pre-RoPE +hook, e.g. Qwen3's QK-norm), and its HF attn loader; nothing here switches on a family. + +The decomposed sites are any per-layer weight matrices (SPEC §1/§3) named torch-style: +`layers.{i}.self_attn.{q,k,v,o}_proj` and `layers.{i}.mlp.{gate,up,down}_proj`, each +with its own C. `GLUDecomposedModel` (an `eqx.Module`) carries the full frozen model — +embedding through every layer to the LM head — as array fields, threaded into the jitted +step as a pytree arg; layers without sites run the plain frozen block. + +q/k/v sites are decomposed BEFORE `_prep_qk`/RoPE/SDPA (the masked site output feeds the +attention math); the o site applies to the attention output. V/U masters are fp32 +keyed per site (`DecompVU`); frozen weights are stored bf16 (SPEC N1) — the trainer +casts for compute. + +Real HF weights load straight from the cached safetensors (no torch dep). +""" + +import json +import re +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Protocol + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P +from jax.typing import DTypeLike +from jaxtyping import Array, Float, Int +from safetensors import safe_open + +from param_decomp.components import ( + DecompVU, + SiteC, + SiteSpec, + dequantize_fp8, + quantize_fp8, + site_out, +) +from param_decomp.losses import kl_per_position +from param_decomp.sharding import assert_divisible +from vendored_jax.llama import apply_rope, causal_sdpa, repeat_kv, rms_norm, rope_cos_sin + + +class GLUArch(Protocol): + """The arch-config surface the shared machinery reads. Family configs satisfy it + structurally — the vendored `LlamaConfig` (llama3 rope-scaling fields on top of + these) and the family-neutral `GLUConfig`.""" + + @property + def vocab_size(self) -> int: ... + @property + def n_layer(self) -> int: ... + @property + def n_head(self) -> int: ... + @property + def n_kv_head(self) -> int: ... + @property + def n_embd(self) -> int: ... + @property + def n_intermediate(self) -> int: ... + @property + def rms_norm_eps(self) -> float: ... + @property + def head_dim(self) -> int: ... + @property + def n_rep(self) -> int: ... + + +@dataclass(frozen=True) +class GLUConfig: + """A family-neutral GLU-transformer arch config (plain RoPE, no family extras). + Families with more knobs bring their own config type satisfying `GLUArch`.""" + + vocab_size: int + n_layer: int + n_head: int + n_kv_head: int + n_embd: int + n_intermediate: int + rope_theta: float + rms_norm_eps: float + max_position_embeddings: int + + @property + def head_dim(self) -> int: + return self.n_embd // self.n_head + + @property + def n_rep(self) -> int: + return self.n_head // self.n_kv_head + + +def default_inv_freq(head_dim: int, rope_theta: float) -> Float[Array, " hd2"]: + """Plain (unscaled) RoPE inverse frequencies — HF's `rope_type: default`.""" + return 1.0 / (rope_theta ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim)) + + +KIND_ORDER = ("q", "k", "v", "o", "gate", "up", "down") +"""Within-layer canonical site order = computation order. The canonical site order +(`glu_site_specs`) is layer-ascending, then this.""" +ATTN_KINDS = ("q", "k", "v", "o") +MLP_KINDS = ("gate", "up", "down") + +SITE_NAME_PATTERN = re.compile( + r"^layers\.(\d+)\.(?:self_attn\.(q|k|v|o)|mlp\.(gate|up|down))_proj$" +) + + +def site_name(layer: int, kind: str) -> str: + assert kind in KIND_ORDER, kind + submodule = "self_attn" if kind in ATTN_KINDS else "mlp" + return f"layers.{layer}.{submodule}.{kind}_proj" + + +def parse_site_name(name: str) -> tuple[int, str]: + """`layers.{i}.{self_attn,mlp}.{kind}_proj` -> (layer, kind); rejects anything else + (including kind/submodule mismatches like `self_attn.gate_proj`).""" + match = SITE_NAME_PATTERN.match(name) + assert match is not None, f"unsupported site name {name!r}" + layer, attn_kind, mlp_kind = match.groups() + return int(layer), attn_kind if attn_kind is not None else mlp_kind + + +def site_dims(cfg: GLUArch, kind: str) -> tuple[int, int]: + """(d_in, d_out) of one per-layer matrix, right-mult orientation.""" + d, di = cfg.n_embd, cfg.n_intermediate + qd = cfg.n_head * cfg.head_dim + kvd = cfg.n_kv_head * cfg.head_dim + match kind: + case "q": + return d, qd + case "k" | "v": + return d, kvd + case "o": + return qd, d + case "gate" | "up": + return d, di + case "down": + return di, d + case _: + raise AssertionError(f"unknown kind {kind!r}") + + +def canonical_site_cs(site_cs: tuple[SiteC, ...]) -> tuple[SiteC, ...]: + """Canonical site order: layer-ascending, `KIND_ORDER` within a layer. Names must + parse and be unique.""" + names = [site.name for site in site_cs] + assert len(set(names)) == len(names), f"duplicate sites in {names}" + + def order_key(site: SiteC) -> tuple[int, int]: + layer, kind = parse_site_name(site.name) + return layer, KIND_ORDER.index(kind) + + return tuple(sorted(site_cs, key=order_key)) + + +def mlp_family_site_cs(first_layer: int, last_layer: int, C: int) -> tuple[SiteC, ...]: + """The gate/up/down sites of a contiguous layer range at one C (the native-config + target family), in canonical order.""" + assert first_layer <= last_layer, (first_layer, last_layer) + return tuple( + SiteC(site_name(layer, kind), C) + for layer in range(first_layer, last_layer + 1) + for kind in MLP_KINDS + ) + + +def glu_site_specs(cfg: GLUArch, site_cs: tuple[SiteC, ...]) -> tuple[SiteSpec, ...]: + """Shape-resolved specs in canonical order (input must already be canonical).""" + assert site_cs == canonical_site_cs(site_cs), f"sites not in canonical order: {site_cs}" + specs = [] + for site in site_cs: + layer, kind = parse_site_name(site.name) + assert 0 <= layer < cfg.n_layer, (site.name, cfg.n_layer) + assert site.C >= 1, site + specs.append(SiteSpec(site.name, *site_dims(cfg, kind), site.C)) + return tuple(specs) + + +# ----------------------------- frozen layers ----------------------------- + + +class FrozenAttn(eqx.Module): + """Plain GQA attention (Llama, LlamaSimpleMLP). A family with extra pre-RoPE math + subclasses and overrides `_prep_qk` (and `shardings` for any extra fields) — e.g. + `qwen3_8b.Qwen3FrozenAttn`'s per-head QK-norm.""" + + wq: Float[Array, "qd d"] + wk: Float[Array, "kvd d"] + wv: Float[Array, "kvd d"] + wo: Float[Array, "d qd"] + n_head: int = eqx.field(static=True) + n_kv_head: int = eqx.field(static=True) + head_dim: int = eqx.field(static=True) + n_rep: int = eqx.field(static=True) + + def _prep_qk( + self, q: Float[Array, "b t h hd"], k: Float[Array, "b t kvh hd"] + ) -> tuple[Array, Array]: + """Family hook between the head reshape and RoPE; identity for plain attention.""" + return q, k + + def shardings(self, mesh: Mesh) -> "FrozenAttn": + """Stacked (leading `n_layer`, UNSHARDED — the scan axis) FSDP on `fsdp`: the `d` dim + shards on `fsdp` (gathered per layer in the scan, on NVLink); the HEAD dim stays + REPLICATED. `core` runs batch-parallel attention (q/k/v constrained batch over the + full mesh, heads replicated — that identical spec is what cuDNN's partitioner + requires), so the projections must come out heads-replicated. Attention weights are + small, so FSDP-on-`fsdp` sharding is plenty.""" + in_fsdp = NamedSharding(mesh, P(None, None, "fsdp")) # qkv [nc, head(repl), d on fsdp] + out_fsdp = NamedSharding(mesh, P(None, "fsdp", None)) # wo [nc, d on fsdp, head(repl)] + for w in (self.wq, self.wk, self.wv): + assert_divisible(w.shape[2], mesh, "fsdp", "FrozenAttn qkv in (d)") + assert_divisible(self.wo.shape[1], mesh, "fsdp", "FrozenAttn out-proj out (d)") + return eqx.tree_at( + lambda a: (a.wq, a.wk, a.wv, a.wo), self, (in_fsdp, in_fsdp, in_fsdp, out_fsdp) + ) + + def core( + self, + q_flat: Float[Array, "b t qd"], + k_flat: Float[Array, "b t kvd"], + v_flat: Float[Array, "b t kvd"], + inv_freq: Array, + ) -> Float[Array, "b t qd"]: + """RoPE + causal SDPA between the q/k/v projections and the o projection — + the seam the decomposed q/k/v site outputs feed into.""" + b, t, _ = q_flat.shape + assert q_flat.shape[-1] == self.n_head * self.head_dim, q_flat.shape + assert k_flat.shape[-1] == self.n_kv_head * self.head_dim, k_flat.shape + assert v_flat.shape[-1] == self.n_kv_head * self.head_dim, v_flat.shape + 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) + q, k = self._prep_qk(q, k) + q = q.transpose(0, 2, 1, 3) + k = k.transpose(0, 2, 1, 3) + v = v_flat.reshape(b, t, self.n_kv_head, self.head_dim).transpose(0, 2, 1, 3) + cos, sin = rope_cos_sin(inv_freq, t, q_flat.dtype) + q, k = apply_rope(q, k, cos, sin) + # Native GQA: do NOT repeat_kv — `dot_product_attention` handles the q-heads:kv-heads + # grouping internally. Repeating k/v to n_head and THEN sharding makes the SPMD + # partitioner derive the repeated-k/v layout from the small n_kv_head source, + # inconsistent with q -> cuDNN "Query, key and value should have same sharding" (forward + # AND its rematerialized backward). Real GQA keeps q (n_head) and k/v (n_kv_head) as + # independent head-parallel tensors with the SAME spec (different head COUNTS is fine). + # cuDNN flash attention's custom partitioner requires q/k/v IDENTICALLY sharded. + # Pure HSDP: pin all three batch-parallel over the FULL mesh, HEADS replicated + # (`P(('replicate','fsdp'), None, None, None)` for `[b, heads, t, hd]`). The identical + # q/k/v spec is exactly what cuDNN's flash partitioner demands; heads-replicated keeps + # q (n_head) and k/v (n_kv_head) consistently sharded (no head TP that would split them + # to different per-rank counts). The q/k/v projection outputs are `d_out`-replicated + # (U FSDP's the C side, not d_out's head); this constraint pins the batch right here. + # Guarded so it's a no-op off-mesh (CPU tests / single device); `run.py` sets the + # global mesh. + if not jax.sharding.get_abstract_mesh().empty: + qkv_spec = jax.sharding.PartitionSpec( + ("replicate", "fsdp"), None, None, None + ) # batch-parallel: q/k/v IDENTICAL spec -> cuDNN happy; flash = no score materialization + q, k, v = (jax.lax.with_sharding_constraint(a, qkv_spec) for a in (q, k, v)) + return causal_sdpa(q, k, v).transpose(0, 2, 1, 3).reshape(b, t, self.n_head * self.head_dim) + + def __call__(self, x: Float[Array, "b t d"], inv_freq: Array) -> Array: + return self.core(x @ self.wq.T, x @ self.wk.T, x @ self.wv.T, inv_freq) @ self.wo.T + + def pattern( + self, q_flat: Float[Array, "b t qd"], k_flat: Float[Array, "b t kvd"], inv_freq: Array + ) -> Float[Array, "b h t t"]: + """Post-softmax causal attention map from flat Q/K projections — the target-owned + recipe behind the attn-patterns eval (`GLUDecomposedModel.attn_pattern`). Same + `_prep_qk`/RoPE/GQA math as `core`; scores in fp32, scaled by `1/√head_dim`, + causal-masked, softmaxed (no score materialization concerns — eval-only).""" + b, t, _ = q_flat.shape + assert q_flat.shape[-1] == self.n_head * self.head_dim, q_flat.shape + assert k_flat.shape[-1] == self.n_kv_head * self.head_dim, k_flat.shape + 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) + q, k = self._prep_qk(q, k) + 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, self.n_rep) + scores = jnp.einsum("bhqd,bhkd->bhqk", q.astype(jnp.float32), k.astype(jnp.float32)) + scores = scores / self.head_dim**0.5 + causal = jnp.triu(jnp.ones((t, t), bool), k=1) + return jax.nn.softmax(jnp.where(causal, -jnp.inf, scores), axis=-1) + + +class GLULayer(eqx.Module): + """One layer's frozen weights — norms, attention, MLP. Decomposed sites read + their frozen target W from here at forward time; layers without sites run the + plain frozen block from the same fields. Weights pass as a runtime arg — never + baked into the HLO as a multi-GB constant.""" + + ln1: Float[Array, " d"] + ln2: Float[Array, " d"] + attn: FrozenAttn + Wg: Float[Array, "di d"] + Wu: Float[Array, "di d"] + Wd: Float[Array, "d di"] + + def shardings(self, mesh: Mesh) -> "GLULayer": + """Stacked FSDP on `fsdp` (no TP): every MLP weight shards its `d`-dim on `fsdp`, + the intermediate dim stays replicated; gathered back per layer inside the scan. + Norms replicate; attn delegates to `FrozenAttn.shardings`.""" + in_fsdp = NamedSharding(mesh, P(None, None, "fsdp")) # Wg/Wu [nc, di(repl), d on fsdp] + out_fsdp = NamedSharding(mesh, P(None, "fsdp", None)) # Wd [nc, d on fsdp, di(repl)] + repl = NamedSharding(mesh, P()) + assert_divisible(self.Wg.shape[2], mesh, "fsdp", "Wg in (d)") + assert_divisible(self.Wd.shape[1], mesh, "fsdp", "Wd out (d)") + return eqx.tree_at( + lambda layer: (layer.ln1, layer.ln2, layer.attn, layer.Wg, layer.Wu, layer.Wd), + self, + (repl, repl, self.attn.shardings(mesh), in_fsdp, in_fsdp, out_fsdp), + ) + + +def _frozen_site_weight(layer: GLULayer, kind: str) -> Array: + match kind: + case "q": + return layer.attn.wq + case "k": + return layer.attn.wk + case "v": + return layer.attn.wv + case "o": + return layer.attn.wo + case "gate": + return layer.Wg + case "up": + return layer.Wu + case "down": + return layer.Wd + case _: + raise AssertionError(f"unknown kind {kind!r}") + + +# ----------------------------- forwards ----------------------------- + + +def _clean_mlp_out(layer: GLULayer, mlp_in: Array) -> Array: + """Frozen target MLP — exactly `W` applied, not the `V@U + (W−V@U)` identity, so + non-live sites carry no V/U gradient and no decomposition rounding (SPEC S2/S3).""" + return (jax.nn.silu(mlp_in @ layer.Wg.T) * (mlp_in @ layer.Wu.T)) @ layer.Wd.T + + +def _stack_layers(layers: list[GLULayer]) -> GLULayer: + """Stack a per-layer `GLULayer` list into one whose array leaves carry a leading + layer axis — the `xs` for a `lax.scan` over the (homogeneous) block stack. Static + fields (attn head counts) ride in the treedef, shared across iterations.""" + return jax.tree.map(lambda *per_layer: jnp.stack(per_layer), *layers) + + +def _tap_layer(key: str) -> int: + """Global block index a `read_activations` key reads at: the block a `resid.{L}` tap + enters, or the block a decomposed site lives in.""" + if key.startswith("resid."): + return int(key.split(".")[1]) + return parse_site_name(key)[0] + + +def _per_kind_dims(components: DecompVU) -> dict[str, tuple[int, int, int]]: + """Per decomposed KIND, the `(d_in, C, d_out)` shared across its layers — asserting + uniformity, the precondition for the layer-`lax.scan` masked forward (it stacks each + kind across layers, so every layer's matrix of that kind must be the same shape).""" + kind_dims: dict[str, tuple[int, int, int]] = {} + for name, (V, U) in components.vu.items(): + kind = parse_site_name(name)[1] + dims = (V.shape[0], V.shape[1], U.shape[1]) + assert kind_dims.setdefault(kind, dims) == dims, ( + f"per-kind dims must be uniform across layers for the scan masked forward: " + f"{kind} {dims} != {kind_dims[kind]}" + ) + return kind_dims + + +def _stack_per_kind_vu(components: DecompVU, n_layers: int) -> dict[str, dict[str, Array]]: + """Per decomposed KIND, the layer-stacked `(V, U)` arrays — the MASK-INDEPENDENT part of + the scan inputs (a leading layer axis, one homogeneous body across layers). Mask/live/ + delta/route are attached per-forward by `_attach_per_kind_masks`; the V/U stack + + `_reconstruct_compute_weights` (the ÷N→÷fsdp cross-node gather) are the same for EVERY + forward in a step, so they are built ONCE via `prepare_compute_weights` and shared. + Per-kind dims (d_in, C, d_out) must be uniform across layers (asserted in `_per_kind_dims`).""" + kind_dims = _per_kind_dims(components) + vu_dt = next(iter(components.vu.values()))[0].dtype + per_kind: dict[str, dict[str, Array]] = {} + for kind, (d_in, C, d_out) in kind_dims.items(): + names = [site_name(layer, kind) for layer in range(n_layers)] + Vs = jnp.stack( + [ + components.vu[n][0] if n in components.vu else jnp.zeros((d_in, C), vu_dt) + for n in names + ] + ) + Us = jnp.stack( + [ + components.vu[n][1] if n in components.vu else jnp.zeros((C, d_out), vu_dt) + for n in names + ] + ) + per_kind[kind] = {"V": Vs, "U": Us} + return per_kind + + +def _attach_per_kind_masks( + prepared: dict[str, dict[str, Array]], + n_layers: int, + leading: tuple[int, ...], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live_set: frozenset[str], + has_delta: bool, +) -> dict[str, dict[str, Array]]: + """Attach the per-forward `(live, mask[, delta][, route])` stacks to the shared, already + stacked + ÷fsdp-reconstructed `prepared` per-kind `(V, U)` weights. Sites absent from + `live` get dummy mask/delta/route (the `cond` frozen branch ignores them); `masks`/ + `delta_masks`/`routes` exist only for live sites (recon builds them per-chunk).""" + # Dummy mask/delta/route shapes match the REAL entries (the source scope sets the leading + # shape: `sc` broadcasts over batch as `(1, T)`, not the full `(B, T)`). + a_mask = next(iter(masks.values())) if masks else None + mask_lead = a_mask.shape[:-1] if a_mask is not None else leading + a_delta = next(iter(delta_masks.values())) if (has_delta and delta_masks) else None + a_route = next(iter(routes.values())) if (routes and len(routes)) else None + + per_kind: dict[str, dict[str, Array]] = {} + for kind, vu_entry in prepared.items(): + C = vu_entry["V"].shape[-1] + mask_dt = a_mask.dtype if a_mask is not None else vu_entry["V"].dtype + names = [site_name(layer, kind) for layer in range(n_layers)] + live_flags = jnp.array([n in live_set for n in names]) + masks_k = jnp.stack( + [masks[n] if n in live_set else jnp.ones((*mask_lead, C), mask_dt) for n in names] + ) + entry: dict[str, Array] = {**vu_entry, "live": live_flags, "mask": masks_k} + if has_delta: + d_shape = a_delta.shape if a_delta is not None else leading + d_dt = a_delta.dtype if a_delta is not None else mask_dt + entry["delta"] = jnp.stack( + [delta_masks[n] if n in live_set else jnp.zeros(d_shape, d_dt) for n in names] + ) + if routes is not None: + r_shape = a_route.shape if a_route is not None else leading + r_dt = a_route.dtype if a_route is not None else jnp.bool_ + entry["route"] = jnp.stack( + [routes[n] if n in live_set else jnp.zeros(r_shape, r_dt) for n in names] + ) + per_kind[kind] = entry + return per_kind + + +def _stack_ci_per_kind(ci_lower: dict[str, Array], n_layers: int) -> dict[str, Array]: + """Stack the per-site CI envelope into per-kind `[n_layer, *leading, C]` — built ONCE per + step and SHARED across every stochastic recon forward (the CI envelope is identical for + all). Mirrors `_stack_per_kind_vu`: the shared stack replaces N per-forward mask stacks.""" + kinds: dict[str, Array] = {} + sample_by_kind: dict[str, Array] = {} + for name, v in ci_lower.items(): + sample_by_kind.setdefault(parse_site_name(name)[1], v) + for kind, sample in sample_by_kind.items(): + names = [site_name(layer, kind) for layer in range(n_layers)] + kinds[kind] = jnp.stack( + [ci_lower[n] if n in ci_lower else jnp.zeros_like(sample) for n in names] + ) + return kinds + + +def _attach_per_kind_stochastic( + prepared: dict[str, dict[str, Array]], + n_layers: int, + leading: tuple[int, ...], + ci_stacked: dict[str, Array], + draw_key: Array, + routes: dict[str, Array] | None, + live_set: frozenset[str], +) -> dict[str, dict[str, Array]]: + """Stochastic recon: attach the SHARED per-kind `ci` stack + per-(layer,kind) RNG keys + instead of pre-built mask/delta stacks. `masked_site` draws `source = uniform(key)` and + builds `mask = ci + (1−ci)·source` INSIDE the checkpointed block, so the per-forward mask + is recomputed in the backward (faithful by checkpoint determinism — same key fwd+bwd) and + never held. Only the (tiny) keys + live-flags are per-forward; the `[n_layer,*,C]` ci stack + is shared, so N forwards' mask stacks collapse to one ci stack (the memory win).""" + src_base, delta_base = jax.random.split(draw_key) + a_route = next(iter(routes.values())) if (routes and len(routes)) else None + per_kind: dict[str, dict[str, Array]] = {} + for kind, vu_entry in prepared.items(): + names = [site_name(layer, kind) for layer in range(n_layers)] + live_flags = jnp.array([n in live_set for n in names]) + kind_idx = KIND_ORDER.index(kind) + src_keys = jnp.stack( + [ + jax.random.fold_in(jax.random.fold_in(src_base, kind_idx), layer) + for layer in range(n_layers) + ] + ) + delta_keys = jnp.stack( + [ + jax.random.fold_in(jax.random.fold_in(delta_base, kind_idx), layer) + for layer in range(n_layers) + ] + ) + entry: dict[str, Array] = { + **vu_entry, "live": live_flags, "ci": ci_stacked[kind], + "src_key": src_keys, "delta_key": delta_keys, + } # fmt: skip + if routes is not None: + r_shape = a_route.shape if a_route is not None else leading + r_dt = a_route.dtype if a_route is not None else jnp.bool_ + entry["route"] = jnp.stack( + [routes[n] if n in live_set else jnp.zeros(r_shape, r_dt) for n in names] + ) + per_kind[kind] = entry + return per_kind + + +def _reconstruct_compute_weights( + per_kind: dict[str, dict[str, Array]], + fp8: bool, +) -> dict[str, dict[str, Array]]: + """The ZeRO-1 weight reconstruction (pure-HSDP backup layout). The stacked + `[n_layer, d_in, C]` / `[n_layer, C, d_out]` compute weights arrive with their FSDP dim + sharded ÷N over the FULL mesh (the master is `P(("replicate","fsdp"), ...)`). Reconstruct + them to the `fsdp`-sharded (÷fsdp) COMPUTE layout here — BEFORE the layer scan — so: + + * the cross-`replicate` gather runs ONCE per step in ENTRY (off the hot path), + landing a SMALL ÷fsdp-resident weight stack (`[n_layer, d_in/fsdp, C]`), NOT the full + `[n_layer, d_in, C]` model; + * the per-layer scan body then gathers ONE layer's `fsdp` shard to full d_in + transiently (intra-node NVLink), freed each iteration — never a full-model resident. + + Cast to bf16 HERE (not f32) so the ÷fsdp-resident compute stack is half-size and XLA + can't keep an f32 full copy alive. The fp32 masters + Adam stay ÷N (untouched — this is + a separate read-only compute view). The leading `n_layer` axis (the scan `xs`) stays + unsharded. No-op off-mesh (CPU / single device); `run.py` sets the global mesh.""" + if jax.sharding.get_abstract_mesh().empty: + return per_kind + v_spec = P(None, "fsdp", "tp") # [n_layer, d_in ÷fsdp, C ÷tp] — d gathered/step, C stays ÷tp + u_spec = P(None, "tp", "fsdp") # [n_layer, C ÷tp, d_out ÷fsdp] + out: dict[str, dict[str, Array]] = {} + for kind, entry in per_kind.items(): + pinned = dict(entry) + # optimization_barrier forces the cast/quant to materialize BEFORE the ÷N→÷fsdp gather, + # so the collective moves the compute dtype — XLA otherwise sinks the convert past the + # all-gather and gathers the f32 master (2x the comm; the convert gathers in the HLO). + if fp8: + # Quantized All-Gather: the ÷fsdp compute weights are fp8; the per-layer ÷fsdp→full + # gather then moves fp8 (½ the bf16 bytes), dequantized to bf16 in `masked_site`. + # Per-tensor scalar scale rides alongside (replicated, survives the gather). + vq, vs = quantize_fp8(entry["V"]) + uq, us = quantize_fp8(entry["U"]) + pinned["V"] = jax.lax.with_sharding_constraint(jax.lax.optimization_barrier(vq), v_spec) + pinned["U"] = jax.lax.with_sharding_constraint(jax.lax.optimization_barrier(uq), u_spec) + pinned["V_scale"], pinned["U_scale"] = vs, us + else: + pinned["V"] = jax.lax.with_sharding_constraint( + jax.lax.optimization_barrier(entry["V"].astype(jnp.bfloat16)), v_spec + ) + pinned["U"] = jax.lax.with_sharding_constraint( + jax.lax.optimization_barrier(entry["U"].astype(jnp.bfloat16)), u_spec + ) + out[kind] = pinned + return out + + +class GLUDecomposedModel(eqx.Module): + """The GLU-transformer `DecomposedModel` (the `lm.py` contract; SPEC §1), shared + across the HF GLU families — a family's identity lives in its `stacked.attn` module + (its `FrozenAttn` variant) and `inv_freq`, never in a switch here. + + Carries the FROZEN full model (embedding, all blocks, final norm, lm_head) as array + fields — so it threads into the jitted step as a pytree arg, its weights traced not + baked. The TRAINABLE V/U (`vu: DecompVU`) is passed to the forward methods explicitly: + separate lifecycle (own optimizer + checkpoint, C-sharded while these weights + replicate), so it is NOT a field here. + + Forward methods take token `inputs` and embed internally. Blocks with no decomposed + site run the plain frozen path — so a subset decomposition just leaves the rest + frozen. + + `sites` / `leading_axes` are static config.""" + + embed: Float[Array, "vocab d"] + stacked: GLULayer # the per-layer weights stacked on a leading layer axis (the scan + # `xs`), stored pre-stacked: a saved jit input, never re-stacked inside a forward. + n_layer: int = eqx.field(static=True) + norm: Float[Array, " d"] + lm_head: Float[Array, "vocab d"] + inv_freq: Float[Array, " hd2"] + sites: tuple[SiteSpec, ...] = eqx.field(static=True) + leading_axes: tuple[str, ...] = eqx.field(static=True) + eps: float = eqx.field(static=True) + scan_unroll: int = eqx.field(static=True, default=1) + """`lax.scan(unroll=)` factor over the block stack (`RuntimeConfig.scan_unroll`); 1 = + plain per-layer scan.""" + gather_fp8: bool = eqx.field(static=True, default=False) + """Quantized all-gather of the ÷fsdp compute V/U (`RuntimeConfig.gather_fp8`).""" + + @property + def site_names(self) -> tuple[str, ...]: + return tuple(s.name for s in self.sites) + + @property + def layers(self) -> list[GLULayer]: + """Per-layer view of `stacked` (slices the leading layer axis). For non-hot + consumers (equivalence harness); the forwards use `stacked`.""" + return [jax.tree.map(lambda a, idx=i: a[idx], self.stacked) for i in range(self.n_layer)] + + def attn_pattern( + self, + q_site: str, + q_flat: Float[Array, "b t qd"], + k_flat: Float[Array, "b t kvd"], + ) -> Float[Array, "b h t t"]: + """Post-softmax causal attention map from a layer's flat Q/K projections — the + target-owned recipe the attn-patterns eval consumes (`attn_patterns_eval`'s + `AttnPatternModel` protocol). Delegates to that LAYER's own attention module, so + family behavior (Qwen3's per-layer QK-norm) comes along for free.""" + layer = parse_site_name(q_site)[0] + attn = jax.tree.map(lambda a, idx=layer: a[idx], self.stacked.attn) + return attn.pattern(q_flat, k_flat, self.inv_freq) + + def shardings(self, mesh: Mesh) -> "GLUDecomposedModel": + """FSDP-on-`fsdp` the per-layer weights (`stacked.shardings` — `d` on `fsdp`, + head/intermediate replicated; the ~14 GB layer bulk shards `/fsdp`, gathered per layer + inside the scan, on NVLink). embed / lm_head / norm / inv_freq REPLICATE — the ~2 GB + embed+head is small and vocab-parallel logits/lookup aren't worth the complexity. (The + old all-replicate justification — "the target is small vs activations" — is stale: at + the full 32-layer model the replicated target + its backward/remat copies dominate the + step's peak, which is what this shards away.)""" + repl = NamedSharding(mesh, P()) + return eqx.tree_at( + lambda m: (m.embed, m.norm, m.lm_head, m.inv_freq, m.stacked), + self, + (repl, repl, repl, repl, self.stacked.shardings(mesh)), + ) + + @staticmethod + def recon_loss_fn(masked_output: Array, clean_output: Array) -> Array: + return kl_per_position(masked_output, clean_output) + + def embed_tokens(self, tokens: Int[Array, "b t"]) -> Float[Array, "b t d"]: + return self.embed[tokens] + + def clean_output(self, inputs: Int[Array, "b t"]) -> Array: + """The all-frozen forward — the recon target (SPEC S3). A `lax.scan` over the block + stack so XLA compiles one block body instead of unrolling all 32 layers (the compile + fix for the full model; the scan reassociates float ops vs an unrolled loop, within + fp32 tolerance).""" + + def block(x: Array, layer: GLULayer) -> tuple[Array, None]: + x = x + layer.attn(rms_norm(x, layer.ln1, self.eps), self.inv_freq) + x = x + _clean_mlp_out(layer, rms_norm(x, layer.ln2, self.eps)) + return x, None + + x, _ = jax.lax.scan(block, self.embed_tokens(inputs), self.stacked) + x = rms_norm(x, self.norm, self.eps) + return x @ self.lm_head.T + + def read_activations( + self, inputs: Int[Array, "b t"], wanted: tuple[str, ...] + ) -> dict[str, Array]: + """Frozen-path activation accessor (CI input side, SPEC S4; harvest's per-site + matrix inputs). + + `wanted` keys are either `resid.{layer}` (residual stream ENTERING that block — the + chunkwise CI fn's `input_names`) or a decomposed SITE NAME (the activation entering + that site's weight on the frozen path: `q/k/v_proj` ← post-LN1 residual, `o_proj` ← + the attention output, `gate/up_proj` ← post-LN2 residual, `down_proj` ← + `silu(gate)·up`). The residual is threaded identically to `clean_output`; the + per-site intermediates come from the same RMSNorm/attn/MLP math. Stops once the last + requested key's block is fully covered (no wasted block compute past it).""" + wanted_set = frozenset(wanted) + last = max(_tap_layer(key) for key in wanted) + taps: dict[str, Array] = {} + x = self.embed_tokens(inputs) + for layer in range(self.n_layer): + block = jax.tree.map(lambda a, li=layer: a[li], self.stacked) + if f"resid.{layer}" in wanted_set: + taps[f"resid.{layer}"] = x + attn = block.attn + h1 = rms_norm(x, block.ln1, self.eps) + attn_y = attn.core(h1 @ attn.wq.T, h1 @ attn.wk.T, h1 @ attn.wv.T, self.inv_freq) + post_attn = x + attn_y @ attn.wo.T + mlp_in = rms_norm(post_attn, block.ln2, self.eps) + down_in = jax.nn.silu(mlp_in @ block.Wg.T) * (mlp_in @ block.Wu.T) + for kind, site_input in ( + ("q", h1), ("k", h1), ("v", h1), ("o", attn_y), + ("gate", mlp_in), ("up", mlp_in), ("down", down_in), + ): # fmt: skip + name = site_name(layer, kind) + if name in wanted_set: + taps[name] = site_input + x = post_attn + down_in @ block.Wd.T + if layer == last: + break + assert set(taps) == wanted_set, (sorted(taps), sorted(wanted)) + return taps + + def _run_masked_forward( + self, + prepared: dict[str, dict[str, Array]], + inputs: Int[Array, "b t"], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + remat: bool, + collect: dict[str, Array] | None, + stochastic: tuple[dict[str, Array], Array] | None = None, + collect_activations: dict[str, Array] | None = None, + ) -> Array: + """The masked decomposed forward shared by `masked_output` / `masked_output_stochastic` + / `masked_site_outputs` / `masked_component_activations` (SPEC §1.3, S2). `live_set` is + static at trace, so the forward runs as `[frozen prefix] → [live block] → [frozen suffix]` + static sub-scans (no per-site `lax.cond`): only the live block carries + gathers V/U. + `live`/`has_delta` are static; a non-None `collect` gathers per-live-site decomposed + OUTPUTS (`(x@V)*m@U + …`, SPEC S31), and a non-None `collect_activations` gathers + per-live-site component ACTIVATIONS `x@V` (`[*leading, C]`, mask-independent — the + pre-mask coefficient the arithmetic CI-grid eval visualizes). Assumes layer-aligned, + contiguous chunks (asserted below). + + `prepared` is the shared, stacked + ÷fsdp-reconstructed per-kind `(V, U)` from + `prepare_compute_weights` (built ONCE per step) — this fn only ATTACHES the per-forward + masks, so the ÷N→÷fsdp cross-node gather is not re-run here (SPEC unchanged; numerics + identical — the reconstruction is mask-independent and the same for every forward).""" + live_set = frozenset(live) + resid = self.embed_tokens(inputs) + leading = resid.shape[:-1] + if stochastic is not None: + ci_stacked, draw_key = stochastic + per_kind = _attach_per_kind_stochastic( + prepared, self.n_layer, leading, ci_stacked, draw_key, routes, live_set + ) + else: + per_kind = _attach_per_kind_masks( + prepared, self.n_layer, leading, masks, delta_masks, routes, live_set, has_delta + ) + decomposed_kinds = frozenset(per_kind) + want_collect = collect is not None + want_collect_acts = collect_activations is not None + + # STATIC liveness. `live_set` is known at trace, so the live/frozen choice per site needs + # NO runtime `lax.cond` — and removing it lets XLA pack + prefetch the V/U gathers (the + # cond was a scheduling/packing barrier). We assume LAYER-ALIGNED, CONTIGUOUS chunks + # (every production plan: `into_groups` with sites_per_chunk % n_decomposed_kinds == 0, + # and `one_chunk`); both are asserted below. The forward is then + # [frozen prefix] → [live block] → [frozen suffix], each a static sub-scan; only the live + # block carries V/U and gathers them. + def layer_is_live(layer: int) -> bool: + flags = {site_name(layer, kind) in live_set for kind in decomposed_kinds} + assert len(flags) == 1, ( + f"layer {layer} is partially live ({flags}); the segmented masked forward assumes " + f"layer-aligned chunks (sites_per_chunk % {len(decomposed_kinds)} == 0)" + ) + return flags.pop() + + live_layers = [layer for layer in range(self.n_layer) if layer_is_live(layer)] + if live_layers: + first_live, last_live = live_layers[0], live_layers[-1] + 1 + assert live_layers == list(range(first_live, last_live)), ( + f"live layers must be contiguous, got {live_layers}" + ) + else: + first_live = last_live = 0 + + def decomp_site(x_in: Array, W: Array, e: dict[str, Array]) -> Array: + v, u = e["V"], e["U"] + if "V_scale" in e: # fp8 QAG: gather the fp8 ÷fsdp weight to full d (½ bytes on the + # wire), THEN dequant to bf16 — the barrier keeps the convert after the gather so + # the collective moves fp8, not bf16. + v = dequantize_fp8( + jax.lax.optimization_barrier( + jax.lax.with_sharding_constraint(v, P(None, "tp")) + ), + e["V_scale"], + ) + u = dequantize_fp8( + jax.lax.optimization_barrier( + jax.lax.with_sharding_constraint(u, P("tp", None)) + ), + e["U_scale"], + ) + if "ci" in e: # stochastic recompute: draw source from the per-layer key and build the + # mask INLINE (recomputed in the backward, not held — the shared `ci` stack + tiny + # key replace the per-forward mask stack). + ci = e["ci"] + source = jax.random.uniform(e["src_key"], ci.shape, dtype=ci.dtype) + mask = ci + (1.0 - ci) * source + delta = ( + jax.random.uniform(e["delta_key"], ci.shape[:-1], dtype=ci.dtype) + if has_delta + else None + ) + return site_out(x_in, v, u, W, mask, delta, e.get("route")) + return site_out(x_in, v, u, W, e["mask"], e.get("delta"), e.get("route")) + + def masked_site( + x_in: Array, kind: str, W: Array, pk: dict[str, dict[str, Array]] + ) -> tuple[Array, Array | None, Array | None]: + # LIVE block only: every decomposed kind decomps (static — no cond); a kind absent + # from the decomposition stays frozen. + if kind not in decomposed_kinds: + return x_in @ W.T, None, None + e = pk[kind] + out = decomp_site(x_in, W, e) + act = ( + (x_in @ e["V"]) if want_collect_acts else None + ) # pre-mask coeff; see masked_component_activations + return out, (out if want_collect else None), act + + def live_block( + x: Array, layer_in: tuple[GLULayer, dict[str, dict[str, Array]]] + ) -> tuple[Array, tuple[dict[str, Array] | None, dict[str, Array] | None]]: + sl, pk = layer_in + attn = sl.attn + h1 = rms_norm(x, sl.ln1, self.eps) + q, qc, qa = masked_site(h1, "q", attn.wq, pk) + k, kc, ka = masked_site(h1, "k", attn.wk, pk) + v, vc, va = masked_site(h1, "v", attn.wv, pk) + attn_y = attn.core(q, k, v, self.inv_freq) + o, oc, oa = masked_site(attn_y, "o", attn.wo, pk) + post_attn = x + o + h2 = rms_norm(post_attn, sl.ln2, self.eps) + g, gc, ga = masked_site(h2, "gate", sl.Wg, pk) + u, uc, ua = masked_site(h2, "up", sl.Wu, pk) + d, dc, da = masked_site(jax.nn.silu(g) * u, "down", sl.Wd, pk) + x = post_attn + d + kinds = ("q", "k", "v", "o", "gate", "up", "down") + collected = ( + { + n: c + for n, c in zip(kinds, (qc, kc, vc, oc, gc, uc, dc), strict=True) + if c is not None + } + if want_collect + else None + ) + collected_acts = ( + { + n: a + for n, a in zip(kinds, (qa, ka, va, oa, ga, ua, da), strict=True) + if a is not None + } + if want_collect_acts + else None + ) + return x, (collected, collected_acts) + + def frozen_block(x: Array, sl: GLULayer) -> tuple[Array, None]: + # Bit-identical to a frozen `masked_site` branch (`x @ Wᵀ` per site), shared with + # `clean_output`. Carries NO V/U → a frozen segment gathers nothing. + x = x + sl.attn(rms_norm(x, sl.ln1, self.eps), self.inv_freq) + x = x + _clean_mlp_out(sl, rms_norm(x, sl.ln2, self.eps)) + return x, None + + # Per-LAYER checkpoint of the scan BODY in BOTH modes — `remat` controls ONLY whether the + # layer ACTIVATIONS are recomputed; it NEVER controls the ÷fsdp→full V/U gather. That + # gather is a NON-dot collective, so it is never a saved residual under either policy — it + # re-gathers in the backward, transient one layer at a time (without checkpoint XLA keeps + # every layer's full gathered V/U live across the scan → OOM). + # remat=True → nothing_saveable: recompute activations AND the gather (min memory). + # remat=False → dots_saveable: SAVE the activation matmuls (no batch dims here → they + # qualify) but still recompute the gather + cheap elementwise. Pure recompute either + # way; zero numerics change. + # `scan_unroll` (native `lax.scan(unroll=k)`) emits k iterations straight-line so XLA can + # prefetch gather(L+1) under matmul(L) — the overlap a 1-layer while-body denies. + policy = ( + jax.checkpoint_policies.nothing_saveable + if remat + else jax.checkpoint_policies.dots_saveable + ) + + def run_scan(body: Any, carry: Array, xs: Any) -> tuple[Array, Any]: + return jax.lax.scan( + jax.checkpoint(body, policy=policy), carry, xs, unroll=self.scan_unroll + ) + + def slice_layers(lo: int, hi: int) -> GLULayer: + return jax.tree.map(lambda a: a[lo:hi], self.stacked) + + x = resid + ys: tuple[dict[str, Array] | None, dict[str, Array] | None] | None = None + if first_live > 0: + x, _ = run_scan(frozen_block, x, slice_layers(0, first_live)) + if last_live > first_live: + pk_live = { + kind: {k: v[first_live:last_live] for k, v in e.items()} + for kind, e in per_kind.items() + } + x, ys = run_scan(live_block, x, (slice_layers(first_live, last_live), pk_live)) + if last_live < self.n_layer: + x, _ = run_scan(frozen_block, x, slice_layers(last_live, self.n_layer)) + + x = rms_norm(x, self.norm, self.eps) + logits = x @ self.lm_head.T + if collect is not None or collect_activations is not None: + assert ys is not None # requested -> the live block emitted the per-kind stacks + ys_out, ys_acts = ys + for sink, stacked in ((collect, ys_out), (collect_activations, ys_acts)): + if sink is None: + continue + assert stacked is not None # sink requested -> the live block emitted this stack + for site in live: + layer, kind = parse_site_name(site) + sink[site] = stacked[kind][layer - first_live] + return logits + + def prepare_compute_weights(self, vu: DecompVU) -> dict[str, dict[str, Array]]: + """Build the shared per-kind compute weights ONCE per step (SPEC unchanged): stack the + per-site V/U into the layer-stacked `[n_layer, …]` form and run the ÷N→÷fsdp cross-node + reconstruction + bf16 cast. The result is mask-independent and identical for every + forward in the step, so the engine builds it once and threads it into all + `masked_output` / `masked_site_outputs` calls — the cross-node gather then runs ONCE per + step (ENTRY) instead of once per forward (the per-forward re-gather was ~10 co-resident + copies of the ÷fsdp stack at peak).""" + return _reconstruct_compute_weights(_stack_per_kind_vu(vu, self.n_layer), self.gather_fp8) + + def masked_output( + self, + prepared: dict[str, dict[str, Array]], + inputs: Int[Array, "b t"], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + *, + remat: bool, + ) -> Array: + return self._run_masked_forward( + prepared, inputs, masks, delta_masks, routes, live, has_delta, remat, None, None + ) + + def stack_ci(self, ci_lower: dict[str, Array]) -> dict[str, Array]: + """Per-kind `[n_layer, *leading, C]` stack of the CI envelope, built ONCE per step and + shared across all stochastic recon forwards (`masked_output_stochastic`). The + StochasticReconCapable capability (SPEC unchanged — pure recompute restructuring).""" + return _stack_ci_per_kind(ci_lower, self.n_layer) + + def masked_output_stochastic( + self, + prepared: dict[str, dict[str, Array]], + inputs: Int[Array, "b t"], + ci_stacked: dict[str, Array], + draw_key: Array, + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + *, + remat: bool, + ) -> Array: + """Stochastic recon forward that RECOMPUTES masks in-block (memory win): the shared + `ci_stacked` + per-layer keys from `draw_key` replace the per-forward mask stack; each + live site draws `source = uniform(key)` and forms `mask = ci + (1−ci)·source` inside the + checkpointed block (faithful by checkpoint determinism). Same forward semantics as + `masked_output` with stochastic sources — only the masks' liverange changes.""" + return self._run_masked_forward( + prepared, inputs, {}, {}, routes, live, has_delta, remat, None, (ci_stacked, draw_key) + ) + + def masked_site_outputs( + self, + prepared: dict[str, dict[str, Array]], + inputs: Int[Array, "b t"], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + ) -> dict[str, Array]: + """Per-`live`-site decomposed output of the masked forward (SPEC S31). Runs the + exact `masked_output` forward, discards the logits, returns the collected outputs.""" + collect: dict[str, Array] = {} + self._run_masked_forward( + prepared, inputs, masks, delta_masks, routes, live, has_delta, False, collect, None + ) + assert set(collect) == set(live), (sorted(collect), sorted(live)) + return collect + + def masked_component_activations( + self, + prepared: dict[str, dict[str, Array]], + inputs: Int[Array, "b t"], + masks: dict[str, Array], + delta_masks: dict[str, Array], + routes: dict[str, Array] | None, + live: tuple[str, ...], + has_delta: bool, + ) -> dict[str, Array]: + """Per-`live`-site component activation `x@V` (`[*leading, C]`, the coefficient + BEFORE the per-component `*mask`) from the masked forward. A site's OWN mask does not + enter its `x@V`, but the masks passed still shape a downstream site's input (e.g. + `down`'s input is the masked `silu(gate)*up`) — so pass the masks the visualization + wants the forward to run under. For the arithmetic CI-grid eval's activation heatmaps; + LM-only, off the recon path.""" + collect_activations: dict[str, Array] = {} + self._run_masked_forward( + prepared, inputs, masks, delta_masks, routes, live, has_delta, False, None, + collect_activations=collect_activations, + ) # fmt: skip + assert set(collect_activations) == set(live), (sorted(collect_activations), sorted(live)) + return collect_activations + + def weight_deltas(self, vu: DecompVU) -> dict[str, Array]: + """fp32 `W − V@U` per site from fp32 masters (SPEC N2; faithfulness input).""" + out: dict[str, Array] = {} + for spec in self.sites: + layer, kind = parse_site_name(spec.name) + W = _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind) + V, U = vu.site(spec.name) + out[spec.name] = ( + W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T + ) + return out + + +# ----------------------------- HF weight loading ----------------------------- + + +def hf_snapshot_dir(model_name: str) -> Path: + """Newest local snapshot of `model_name`. `HF_HUB_CACHE` overrides; otherwise on + cluster (`DATA_MOUNT` set) the shared world-readable cache is the source — a home + `~/.cache` hub is silently mutable, and a wiped entry strands running jobs that + reload weights on requeue.""" + import os + + default_cache = ( + f"{os.environ['DATA_MOUNT']}/artifacts/hf_cache/hub" + if "DATA_MOUNT" in os.environ + else str(Path.home() / ".cache/huggingface/hub") + ) + cache = Path(os.environ.get("HF_HUB_CACHE", default_cache)) + repo = "models--" + model_name.replace("/", "--") + snaps = sorted((cache / repo / "snapshots").iterdir()) + assert snaps, f"no snapshot for {model_name} under {cache}" + return snaps[-1] + + +class HFWeights: + """Lazy keyed access to the sharded safetensors of an HF checkpoint, cast to the + FAMILY's frozen-weights dtype (SPEC N1: bf16 storage — the families pass it; this + module holds no dtype opinion).""" + + def __init__(self, snapshot: Path, dtype: DTypeLike): + index = json.loads((snapshot / "model.safetensors.index.json").read_text()) + self._key_to_file = index["weight_map"] + self._snapshot = snapshot + self._dtype = dtype + self._open: dict[str, Any] = {} + + def get(self, key: str) -> Array: + fname = self._key_to_file[key] + if fname not in self._open: + self._open[fname] = safe_open(str(self._snapshot / fname), framework="numpy") + return jnp.asarray(np.array(self._open[fname].get_tensor(key)), dtype=self._dtype) + + +AttnLoader = Callable[[HFWeights, int], FrozenAttn] +"""`(weights, layer_idx) -> the family's FrozenAttn` — each family file supplies its own +(Llama: plain projections; Qwen3: plus the q/k_norm keys).""" + + +def load_glu_blocks(w: HFWeights, cfg: GLUArch, load_attn: AttnLoader) -> list[GLULayer]: + pre = "model.layers" + return [ + GLULayer( + ln1=w.get(f"{pre}.{i}.input_layernorm.weight"), + ln2=w.get(f"{pre}.{i}.post_attention_layernorm.weight"), + attn=load_attn(w, i), + Wg=w.get(f"{pre}.{i}.mlp.gate_proj.weight"), + Wu=w.get(f"{pre}.{i}.mlp.up_proj.weight"), + Wd=w.get(f"{pre}.{i}.mlp.down_proj.weight"), + ) + for i in range(cfg.n_layer) + ] + + +def build_decomposed_lm( + embed: Array, + layers: list[GLULayer], + norm: Array, + lm_head: Array, + inv_freq: Array, + cfg: GLUArch, + sites: tuple[SiteSpec, ...], + scan_unroll: int = 1, + gather_fp8: bool = False, +) -> GLUDecomposedModel: + """Assemble a `GLUDecomposedModel` from the frozen full-model arrays + decomposition + config. `sites` must be canonical-ordered with dims matching `cfg`. `scan_unroll` / + `gather_fp8` are the `RuntimeConfig` compute knobs (1 / off = the default forward).""" + site_cs = tuple(SiteC(s.name, s.C) for s in sites) + assert sites == glu_site_specs(cfg, canonical_site_cs(site_cs)), ( + f"sites are not the canonical specs for this config: {sites}" + ) + return GLUDecomposedModel( + embed=embed, + stacked=_stack_layers(layers), + n_layer=len(layers), + norm=norm, + lm_head=lm_head, + inv_freq=inv_freq, + sites=sites, + leading_axes=("sequence",), + eps=cfg.rms_norm_eps, + scan_unroll=scan_unroll, + gather_fp8=gather_fp8, + ) + + +def load_decomposed_glu_from_hf( + model_name: str, + cfg: GLUArch, + sites: tuple[SiteSpec, ...], + load_attn: AttnLoader, + inv_freq: Array, + weights_dtype: DTypeLike, + scan_unroll: int = 1, + gather_fp8: bool = False, +) -> GLUDecomposedModel: + """Load a GLU-transformer `DecomposedModel` from the cached HF snapshot: the full + frozen model (embedding, all blocks, final norm, lm_head) as fields plus the static + decomposition config (`sites`). The FAMILY contributes `load_attn` + `inv_freq` (its + RoPE flavor); blocks without a decomposed site run the plain frozen path. + `scan_unroll` / `gather_fp8` are the `RuntimeConfig` compute knobs.""" + w = HFWeights(hf_snapshot_dir(model_name), weights_dtype) + return build_decomposed_lm( + embed=w.get("model.embed_tokens.weight"), + layers=load_glu_blocks(w, cfg, load_attn), + norm=w.get("model.norm.weight"), + lm_head=w.get("lm_head.weight"), + inv_freq=inv_freq, + cfg=cfg, + sites=sites, + scan_unroll=scan_unroll, + gather_fp8=gather_fp8, + ) diff --git a/param_decomp/targets/llama8b_sharding.py b/param_decomp/targets/glu_transformer_sharding.py similarity index 95% rename from param_decomp/targets/llama8b_sharding.py rename to param_decomp/targets/glu_transformer_sharding.py index ad446c04d..386e341a8 100644 --- a/param_decomp/targets/llama8b_sharding.py +++ b/param_decomp/targets/glu_transformer_sharding.py @@ -1,4 +1,4 @@ -"""GSPMD sharding plan for the Llama-8B single-pool step — the pure-HSDP memory story. +"""GSPMD sharding plan for the GLU-transformer (Llama-8B / Qwen3-8B) single-pool step — the pure-HSDP memory story. The 2-D `(replicate, fsdp)` mesh: `fsdp` is the 8 intra-node NVLink GPUs (the FSDP weight-gather / grad-reduce axis), `replicate` the across-node axis. There is NO TP / @@ -12,7 +12,7 @@ fp32 Adam m/v are the dominant non-activation footprint; true ZeRO-1 ÷N (master + m + v each ÷(replicate·fsdp)) takes them from ÷fsdp (≈76 GB/GPU fixed) to ÷N (≈5 GB, scaling). COMPUTE re-pins the bf16 weights to `fsdp`-only ONCE per step (the ZeRO-1 reconstruction, - in ENTRY, off the per-layer hot path; see `llama8b._reconstruct_compute_weights`). + in ENTRY, off the per-layer hot path; see `glu_transformer._reconstruct_compute_weights`). * CI fn + Adam states: sharded ÷N over the full mesh along d_model (in_proj / blocks / heads), same ZeRO-1 reconstruction to `fsdp`-only before the chunk scan. * PGD source (broadcast scope, `{site: (1,T,C+1)}`): REPLICATED. A single adversarial @@ -66,7 +66,7 @@ from param_decomp.configs import BSCScope, SCScope from param_decomp.sharding import hsdp_mesh, place_via_shardings from param_decomp.sharding import shard_batch as _generic_shard_batch -from param_decomp.targets.llama8b import LlamaDecomposedModel +from param_decomp.targets.glu_transformer import GLUDecomposedModel __all__ = [ "hsdp_mesh", @@ -78,7 +78,7 @@ ] -def place_target(tgt: LlamaDecomposedModel, mesh: Mesh) -> LlamaDecomposedModel: +def place_target(tgt: GLUDecomposedModel, mesh: Mesh) -> GLUDecomposedModel: """Eager `device_put` of the already-loaded frozen target onto its own declared placement (`tgt.shardings(mesh)` — FSDP-on-`fsdp`).""" return place_via_shardings(tgt, tgt.shardings(mesh)) diff --git a/param_decomp/targets/llama8b.py b/param_decomp/targets/llama8b.py index 1849df0f3..a4fa8bd5a 100644 --- a/param_decomp/targets/llama8b.py +++ b/param_decomp/targets/llama8b.py @@ -1,63 +1,17 @@ -"""Llama-3.1-8B vendored target — the first `DecomposedModel` implementation. +"""The Llama-3.1-8B family target: its arch config (the vendored `LlamaConfig`, llama3 +rope scaling) and its HF loader over the shared GLU-transformer machinery +(`targets/glu_transformer.py` — plain `FrozenAttn`, no pre-RoPE extras).""" -The decomposed sites are any per-layer weight matrices (SPEC §1/§3) named torch-style: -`layers.{i}.self_attn.{q,k,v,o}_proj` and `layers.{i}.mlp.{gate,up,down}_proj`, each -with its own C. `LlamaDecomposedModel` (an `eqx.Module`) carries the full frozen model — -embedding through every layer to the LM head — as array fields, threaded into the jitted -step as a pytree arg; layers without sites run the plain frozen block. - -q/k/v sites are decomposed BEFORE RoPE/SDPA (the masked site output feeds the -attention math); the o site applies to the attention output. V/U masters are fp32 -keyed per site (`DecompVU`); frozen weights are stored bf16 (SPEC N1) — the trainer -casts for compute. - -Real HF weights load straight from the cached safetensors (no torch dep). -""" - -import json -import re -from pathlib import Path -from typing import Any - -import equinox as eqx -import jax import jax.numpy as jnp -import numpy as np -from jax.sharding import Mesh, NamedSharding -from jax.sharding import PartitionSpec as P -from jaxtyping import Array, Float, Int -from safetensors import safe_open -from param_decomp.components import ( - DecompVU, - SiteC, - SiteSpec, - dequantize_fp8, - quantize_fp8, - site_out, -) -from param_decomp.losses import kl_per_position -from param_decomp.sharding import assert_divisible -from vendored_jax.llama import ( - LlamaConfig, - apply_rope, - causal_sdpa, - llama3_inv_freq, - rms_norm, - rope_cos_sin, -) - -DT = jnp.bfloat16 - -KIND_ORDER = ("q", "k", "v", "o", "gate", "up", "down") -"""Within-layer canonical site order = computation order. The canonical site order -(`llama_site_specs`) is layer-ascending, then this.""" -ATTN_KINDS = ("q", "k", "v", "o") -MLP_KINDS = ("gate", "up", "down") - -SITE_NAME_PATTERN = re.compile( - r"^layers\.(\d+)\.(?:self_attn\.(q|k|v|o)|mlp\.(gate|up|down))_proj$" +from param_decomp.components import SiteSpec +from param_decomp.targets.glu_transformer import ( + FrozenAttn, + GLUDecomposedModel, + HFWeights, + load_decomposed_glu_from_hf, ) +from vendored_jax.llama import LlamaConfig, llama3_inv_freq def llama31_8b_config() -> LlamaConfig: @@ -78,902 +32,7 @@ def llama31_8b_config() -> LlamaConfig: ) -def site_name(layer: int, kind: str) -> str: - assert kind in KIND_ORDER, kind - submodule = "self_attn" if kind in ATTN_KINDS else "mlp" - return f"layers.{layer}.{submodule}.{kind}_proj" - - -def parse_site_name(name: str) -> tuple[int, str]: - """`layers.{i}.{self_attn,mlp}.{kind}_proj` -> (layer, kind); rejects anything else - (including kind/submodule mismatches like `self_attn.gate_proj`).""" - match = SITE_NAME_PATTERN.match(name) - assert match is not None, f"unsupported site name {name!r}" - layer, attn_kind, mlp_kind = match.groups() - return int(layer), attn_kind if attn_kind is not None else mlp_kind - - -def site_dims(cfg: LlamaConfig, kind: str) -> tuple[int, int]: - """(d_in, d_out) of one per-layer matrix, right-mult orientation.""" - d, di = cfg.n_embd, cfg.n_intermediate - qd = cfg.n_head * cfg.head_dim - kvd = cfg.n_kv_head * cfg.head_dim - match kind: - case "q": - return d, qd - case "k" | "v": - return d, kvd - case "o": - return qd, d - case "gate" | "up": - return d, di - case "down": - return di, d - case _: - raise AssertionError(f"unknown kind {kind!r}") - - -def canonical_site_cs(site_cs: tuple[SiteC, ...]) -> tuple[SiteC, ...]: - """Canonical site order: layer-ascending, `KIND_ORDER` within a layer. Names must - parse and be unique.""" - names = [site.name for site in site_cs] - assert len(set(names)) == len(names), f"duplicate sites in {names}" - - def order_key(site: SiteC) -> tuple[int, int]: - layer, kind = parse_site_name(site.name) - return layer, KIND_ORDER.index(kind) - - return tuple(sorted(site_cs, key=order_key)) - - -def mlp_family_site_cs(first_layer: int, last_layer: int, C: int) -> tuple[SiteC, ...]: - """The gate/up/down sites of a contiguous layer range at one C (the native-config - target family), in canonical order.""" - assert first_layer <= last_layer, (first_layer, last_layer) - return tuple( - SiteC(site_name(layer, kind), C) - for layer in range(first_layer, last_layer + 1) - for kind in MLP_KINDS - ) - - -def llama_site_specs(cfg: LlamaConfig, site_cs: tuple[SiteC, ...]) -> tuple[SiteSpec, ...]: - """Shape-resolved specs in canonical order (input must already be canonical).""" - assert site_cs == canonical_site_cs(site_cs), f"sites not in canonical order: {site_cs}" - specs = [] - for site in site_cs: - layer, kind = parse_site_name(site.name) - assert 0 <= layer < cfg.n_layer, (site.name, cfg.n_layer) - assert site.C >= 1, site - specs.append(SiteSpec(site.name, *site_dims(cfg, kind), site.C)) - return tuple(specs) - - -# ----------------------------- frozen layers ----------------------------- - - -class FrozenAttn(eqx.Module): - wq: Float[Array, "qd d"] - wk: Float[Array, "kvd d"] - wv: Float[Array, "kvd d"] - wo: Float[Array, "d qd"] - n_head: int = eqx.field(static=True) - n_kv_head: int = eqx.field(static=True) - head_dim: int = eqx.field(static=True) - n_rep: int = eqx.field(static=True) - - def shardings(self, mesh: "Mesh") -> "FrozenAttn": - """Stacked (leading `n_layer`, UNSHARDED — the scan axis) FSDP on `fsdp`: the `d` dim - shards on `fsdp` (gathered per layer in the scan, on NVLink); the HEAD dim stays - REPLICATED. `core` runs batch-parallel attention (q/k/v constrained batch over the - full mesh, heads replicated — that identical spec is what cuDNN's partitioner - requires), so the projections must come out heads-replicated. Attention weights are - small, so FSDP-on-`fsdp` sharding is plenty.""" - in_fsdp = NamedSharding(mesh, P(None, None, "fsdp")) # qkv [nc, head(repl), d on fsdp] - out_fsdp = NamedSharding(mesh, P(None, "fsdp", None)) # wo [nc, d on fsdp, head(repl)] - for w in (self.wq, self.wk, self.wv): - assert_divisible(w.shape[2], mesh, "fsdp", "FrozenAttn qkv in (d)") - assert_divisible(self.wo.shape[1], mesh, "fsdp", "FrozenAttn out-proj out (d)") - return eqx.tree_at( - lambda a: (a.wq, a.wk, a.wv, a.wo), self, (in_fsdp, in_fsdp, in_fsdp, out_fsdp) - ) - - def core( - self, - q_flat: Float[Array, "b t qd"], - k_flat: Float[Array, "b t kvd"], - v_flat: Float[Array, "b t kvd"], - inv_freq: Array, - ) -> Float[Array, "b t qd"]: - """RoPE + causal SDPA between the q/k/v projections and the o projection — - the seam the decomposed q/k/v site outputs feed into.""" - b, t, _ = q_flat.shape - assert q_flat.shape[-1] == self.n_head * self.head_dim, q_flat.shape - assert k_flat.shape[-1] == self.n_kv_head * self.head_dim, k_flat.shape - assert v_flat.shape[-1] == self.n_kv_head * self.head_dim, v_flat.shape - q = q_flat.reshape(b, t, self.n_head, self.head_dim).transpose(0, 2, 1, 3) - k = k_flat.reshape(b, t, self.n_kv_head, self.head_dim).transpose(0, 2, 1, 3) - v = v_flat.reshape(b, t, self.n_kv_head, self.head_dim).transpose(0, 2, 1, 3) - cos, sin = rope_cos_sin(inv_freq, t, q_flat.dtype) - q, k = apply_rope(q, k, cos, sin) - # Native GQA: do NOT repeat_kv — `dot_product_attention` handles the q-heads:kv-heads - # grouping internally. Repeating k/v to n_head and THEN sharding makes the SPMD - # partitioner derive the repeated-k/v layout from the small n_kv_head source, - # inconsistent with q -> cuDNN "Query, key and value should have same sharding" (forward - # AND its rematerialized backward). Real GQA keeps q (n_head) and k/v (n_kv_head) as - # independent head-parallel tensors with the SAME spec (different head COUNTS is fine). - # cuDNN flash attention's custom partitioner requires q/k/v IDENTICALLY sharded. - # Pure HSDP: pin all three batch-parallel over the FULL mesh, HEADS replicated - # (`P(('replicate','fsdp'), None, None, None)` for `[b, heads, t, hd]`). The identical - # q/k/v spec is exactly what cuDNN's flash partitioner demands; heads-replicated keeps - # q (n_head) and k/v (n_kv_head) consistently sharded (no head TP that would split them - # to different per-rank counts). The q/k/v projection outputs are `d_out`-replicated - # (U FSDP's the C side, not d_out's head); this constraint pins the batch right here. - # Guarded so it's a no-op off-mesh (CPU tests / single device); `run.py` sets the - # global mesh. - if not jax.sharding.get_abstract_mesh().empty: - qkv_spec = jax.sharding.PartitionSpec( - ("replicate", "fsdp"), None, None, None - ) # batch-parallel: q/k/v IDENTICAL spec -> cuDNN happy; flash = no score materialization - q, k, v = (jax.lax.with_sharding_constraint(a, qkv_spec) for a in (q, k, v)) - return causal_sdpa(q, k, v).transpose(0, 2, 1, 3).reshape(b, t, self.n_head * self.head_dim) - - def __call__(self, x: Float[Array, "b t d"], inv_freq: Array) -> Array: - return self.core(x @ self.wq.T, x @ self.wk.T, x @ self.wv.T, inv_freq) @ self.wo.T - - -class LlamaLayer(eqx.Module): - """One layer's frozen weights — norms, attention, MLP. Decomposed sites read - their frozen target W from here at forward time; layers without sites run the - plain frozen block from the same fields. Weights pass as a runtime arg — never - baked into the HLO as a multi-GB constant.""" - - ln1: Float[Array, " d"] - ln2: Float[Array, " d"] - attn: FrozenAttn - Wg: Float[Array, "di d"] - Wu: Float[Array, "di d"] - Wd: Float[Array, "d di"] - - def shardings(self, mesh: "Mesh") -> "LlamaLayer": - """Stacked FSDP on `fsdp` (no TP): every MLP weight shards its `d`-dim on `fsdp`, - the intermediate dim stays replicated; gathered back per layer inside the scan. - Norms replicate; attn delegates to `FrozenAttn.shardings`.""" - in_fsdp = NamedSharding(mesh, P(None, None, "fsdp")) # Wg/Wu [nc, di(repl), d on fsdp] - out_fsdp = NamedSharding(mesh, P(None, "fsdp", None)) # Wd [nc, d on fsdp, di(repl)] - repl = NamedSharding(mesh, P()) - assert_divisible(self.Wg.shape[2], mesh, "fsdp", "Wg in (d)") - assert_divisible(self.Wd.shape[1], mesh, "fsdp", "Wd out (d)") - return eqx.tree_at( - lambda layer: (layer.ln1, layer.ln2, layer.attn, layer.Wg, layer.Wu, layer.Wd), - self, - (repl, repl, self.attn.shardings(mesh), in_fsdp, in_fsdp, out_fsdp), - ) - - -def _frozen_site_weight(layer: LlamaLayer, kind: str) -> Array: - match kind: - case "q": - return layer.attn.wq - case "k": - return layer.attn.wk - case "v": - return layer.attn.wv - case "o": - return layer.attn.wo - case "gate": - return layer.Wg - case "up": - return layer.Wu - case "down": - return layer.Wd - case _: - raise AssertionError(f"unknown kind {kind!r}") - - -# ----------------------------- forwards ----------------------------- - - -def _clean_mlp_out(layer: LlamaLayer, mlp_in: Array) -> Array: - """Frozen target MLP — exactly `W` applied, not the `V@U + (W−V@U)` identity, so - non-live sites carry no V/U gradient and no decomposition rounding (SPEC S2/S3).""" - return (jax.nn.silu(mlp_in @ layer.Wg.T) * (mlp_in @ layer.Wu.T)) @ layer.Wd.T - - -def _stack_layers(layers: list[LlamaLayer]) -> LlamaLayer: - """Stack a per-layer `LlamaLayer` list into one whose array leaves carry a leading - layer axis — the `xs` for a `lax.scan` over the (homogeneous) block stack. Static - fields (attn head counts) ride in the treedef, shared across iterations.""" - return jax.tree.map(lambda *per_layer: jnp.stack(per_layer), *layers) - - -def _tap_layer(key: str) -> int: - """Global block index a `read_activations` key reads at: the block a `resid.{L}` tap - enters, or the block a decomposed site lives in.""" - if key.startswith("resid."): - return int(key.split(".")[1]) - return parse_site_name(key)[0] - - -def _per_kind_dims(components: DecompVU) -> dict[str, tuple[int, int, int]]: - """Per decomposed KIND, the `(d_in, C, d_out)` shared across its layers — asserting - uniformity, the precondition for the layer-`lax.scan` masked forward (it stacks each - kind across layers, so every layer's matrix of that kind must be the same shape).""" - kind_dims: dict[str, tuple[int, int, int]] = {} - for name, (V, U) in components.vu.items(): - kind = parse_site_name(name)[1] - dims = (V.shape[0], V.shape[1], U.shape[1]) - assert kind_dims.setdefault(kind, dims) == dims, ( - f"per-kind dims must be uniform across layers for the scan masked forward: " - f"{kind} {dims} != {kind_dims[kind]}" - ) - return kind_dims - - -def _stack_per_kind_vu(components: DecompVU, n_layers: int) -> dict[str, dict[str, Array]]: - """Per decomposed KIND, the layer-stacked `(V, U)` arrays — the MASK-INDEPENDENT part of - the scan inputs (a leading layer axis, one homogeneous body across layers). Mask/live/ - delta/route are attached per-forward by `_attach_per_kind_masks`; the V/U stack + - `_reconstruct_compute_weights` (the ÷N→÷fsdp cross-node gather) are the same for EVERY - forward in a step, so they are built ONCE via `prepare_compute_weights` and shared. - Per-kind dims (d_in, C, d_out) must be uniform across layers (asserted in `_per_kind_dims`).""" - kind_dims = _per_kind_dims(components) - vu_dt = next(iter(components.vu.values()))[0].dtype - per_kind: dict[str, dict[str, Array]] = {} - for kind, (d_in, C, d_out) in kind_dims.items(): - names = [site_name(layer, kind) for layer in range(n_layers)] - Vs = jnp.stack( - [ - components.vu[n][0] if n in components.vu else jnp.zeros((d_in, C), vu_dt) - for n in names - ] - ) - Us = jnp.stack( - [ - components.vu[n][1] if n in components.vu else jnp.zeros((C, d_out), vu_dt) - for n in names - ] - ) - per_kind[kind] = {"V": Vs, "U": Us} - return per_kind - - -def _attach_per_kind_masks( - prepared: dict[str, dict[str, Array]], - n_layers: int, - leading: tuple[int, ...], - masks: dict[str, Array], - delta_masks: dict[str, Array], - routes: dict[str, Array] | None, - live_set: frozenset[str], - has_delta: bool, -) -> dict[str, dict[str, Array]]: - """Attach the per-forward `(live, mask[, delta][, route])` stacks to the shared, already - stacked + ÷fsdp-reconstructed `prepared` per-kind `(V, U)` weights. Sites absent from - `live` get dummy mask/delta/route (the `cond` frozen branch ignores them); `masks`/ - `delta_masks`/`routes` exist only for live sites (recon builds them per-chunk).""" - # Dummy mask/delta/route shapes match the REAL entries (the source scope sets the leading - # shape: `sc` broadcasts over batch as `(1, T)`, not the full `(B, T)`). - a_mask = next(iter(masks.values())) if masks else None - mask_lead = a_mask.shape[:-1] if a_mask is not None else leading - a_delta = next(iter(delta_masks.values())) if (has_delta and delta_masks) else None - a_route = next(iter(routes.values())) if (routes and len(routes)) else None - - per_kind: dict[str, dict[str, Array]] = {} - for kind, vu_entry in prepared.items(): - C = vu_entry["V"].shape[-1] - mask_dt = a_mask.dtype if a_mask is not None else vu_entry["V"].dtype - names = [site_name(layer, kind) for layer in range(n_layers)] - live_flags = jnp.array([n in live_set for n in names]) - masks_k = jnp.stack( - [masks[n] if n in live_set else jnp.ones((*mask_lead, C), mask_dt) for n in names] - ) - entry: dict[str, Array] = {**vu_entry, "live": live_flags, "mask": masks_k} - if has_delta: - d_shape = a_delta.shape if a_delta is not None else leading - d_dt = a_delta.dtype if a_delta is not None else mask_dt - entry["delta"] = jnp.stack( - [delta_masks[n] if n in live_set else jnp.zeros(d_shape, d_dt) for n in names] - ) - if routes is not None: - r_shape = a_route.shape if a_route is not None else leading - r_dt = a_route.dtype if a_route is not None else jnp.bool_ - entry["route"] = jnp.stack( - [routes[n] if n in live_set else jnp.zeros(r_shape, r_dt) for n in names] - ) - per_kind[kind] = entry - return per_kind - - -def _stack_ci_per_kind(ci_lower: dict[str, Array], n_layers: int) -> dict[str, Array]: - """Stack the per-site CI envelope into per-kind `[n_layer, *leading, C]` — built ONCE per - step and SHARED across every stochastic recon forward (the CI envelope is identical for - all). Mirrors `_stack_per_kind_vu`: the shared stack replaces N per-forward mask stacks.""" - kinds: dict[str, Array] = {} - sample_by_kind: dict[str, Array] = {} - for name, v in ci_lower.items(): - sample_by_kind.setdefault(parse_site_name(name)[1], v) - for kind, sample in sample_by_kind.items(): - names = [site_name(layer, kind) for layer in range(n_layers)] - kinds[kind] = jnp.stack( - [ci_lower[n] if n in ci_lower else jnp.zeros_like(sample) for n in names] - ) - return kinds - - -def _attach_per_kind_stochastic( - prepared: dict[str, dict[str, Array]], - n_layers: int, - leading: tuple[int, ...], - ci_stacked: dict[str, Array], - draw_key: Array, - routes: dict[str, Array] | None, - live_set: frozenset[str], -) -> dict[str, dict[str, Array]]: - """Stochastic recon: attach the SHARED per-kind `ci` stack + per-(layer,kind) RNG keys - instead of pre-built mask/delta stacks. `masked_site` draws `source = uniform(key)` and - builds `mask = ci + (1−ci)·source` INSIDE the checkpointed block, so the per-forward mask - is recomputed in the backward (faithful by checkpoint determinism — same key fwd+bwd) and - never held. Only the (tiny) keys + live-flags are per-forward; the `[n_layer,*,C]` ci stack - is shared, so N forwards' mask stacks collapse to one ci stack (the memory win).""" - src_base, delta_base = jax.random.split(draw_key) - a_route = next(iter(routes.values())) if (routes and len(routes)) else None - per_kind: dict[str, dict[str, Array]] = {} - for kind, vu_entry in prepared.items(): - names = [site_name(layer, kind) for layer in range(n_layers)] - live_flags = jnp.array([n in live_set for n in names]) - kind_idx = KIND_ORDER.index(kind) - src_keys = jnp.stack( - [ - jax.random.fold_in(jax.random.fold_in(src_base, kind_idx), layer) - for layer in range(n_layers) - ] - ) - delta_keys = jnp.stack( - [ - jax.random.fold_in(jax.random.fold_in(delta_base, kind_idx), layer) - for layer in range(n_layers) - ] - ) - entry: dict[str, Array] = { - **vu_entry, "live": live_flags, "ci": ci_stacked[kind], - "src_key": src_keys, "delta_key": delta_keys, - } # fmt: skip - if routes is not None: - r_shape = a_route.shape if a_route is not None else leading - r_dt = a_route.dtype if a_route is not None else jnp.bool_ - entry["route"] = jnp.stack( - [routes[n] if n in live_set else jnp.zeros(r_shape, r_dt) for n in names] - ) - per_kind[kind] = entry - return per_kind - - -def _reconstruct_compute_weights( - per_kind: dict[str, dict[str, Array]], - fp8: bool, -) -> dict[str, dict[str, Array]]: - """The ZeRO-1 weight reconstruction (pure-HSDP backup layout). The stacked - `[n_layer, d_in, C]` / `[n_layer, C, d_out]` compute weights arrive with their FSDP dim - sharded ÷N over the FULL mesh (the master is `P(("replicate","fsdp"), ...)`). Reconstruct - them to the `fsdp`-sharded (÷fsdp) COMPUTE layout here — BEFORE the layer scan — so: - - * the cross-`replicate` gather runs ONCE per step in ENTRY (off the hot path), - landing a SMALL ÷fsdp-resident weight stack (`[n_layer, d_in/fsdp, C]`), NOT the full - `[n_layer, d_in, C]` model; - * the per-layer scan body then gathers ONE layer's `fsdp` shard to full d_in - transiently (intra-node NVLink), freed each iteration — never a full-model resident. - - Cast to bf16 HERE (not f32) so the ÷fsdp-resident compute stack is half-size and XLA - can't keep an f32 full copy alive. The fp32 masters + Adam stay ÷N (untouched — this is - a separate read-only compute view). The leading `n_layer` axis (the scan `xs`) stays - unsharded. No-op off-mesh (CPU / single device); `run.py` sets the global mesh.""" - if jax.sharding.get_abstract_mesh().empty: - return per_kind - v_spec = P(None, "fsdp", "tp") # [n_layer, d_in ÷fsdp, C ÷tp] — d gathered/step, C stays ÷tp - u_spec = P(None, "tp", "fsdp") # [n_layer, C ÷tp, d_out ÷fsdp] - out: dict[str, dict[str, Array]] = {} - for kind, entry in per_kind.items(): - pinned = dict(entry) - # optimization_barrier forces the cast/quant to materialize BEFORE the ÷N→÷fsdp gather, - # so the collective moves the compute dtype — XLA otherwise sinks the convert past the - # all-gather and gathers the f32 master (2x the comm; the convert gathers in the HLO). - if fp8: - # Quantized All-Gather: the ÷fsdp compute weights are fp8; the per-layer ÷fsdp→full - # gather then moves fp8 (½ the bf16 bytes), dequantized to bf16 in `masked_site`. - # Per-tensor scalar scale rides alongside (replicated, survives the gather). - vq, vs = quantize_fp8(entry["V"]) - uq, us = quantize_fp8(entry["U"]) - pinned["V"] = jax.lax.with_sharding_constraint(jax.lax.optimization_barrier(vq), v_spec) - pinned["U"] = jax.lax.with_sharding_constraint(jax.lax.optimization_barrier(uq), u_spec) - pinned["V_scale"], pinned["U_scale"] = vs, us - else: - pinned["V"] = jax.lax.with_sharding_constraint( - jax.lax.optimization_barrier(entry["V"].astype(jnp.bfloat16)), v_spec - ) - pinned["U"] = jax.lax.with_sharding_constraint( - jax.lax.optimization_barrier(entry["U"].astype(jnp.bfloat16)), u_spec - ) - out[kind] = pinned - return out - - -class LlamaDecomposedModel(eqx.Module): - """The Llama-8B `DecomposedModel` (the `lm.py` contract; SPEC §1). - - Carries the FROZEN full model (embedding, all blocks, final norm, lm_head) as array - fields — so it threads into the jitted step as a pytree arg, its weights traced not - baked. The TRAINABLE V/U (`vu: DecompVU`) is passed to the forward methods explicitly: - separate lifecycle (own optimizer + checkpoint, C-sharded while these weights - replicate), so it is NOT a field here. - - Forward methods take token `inputs` and embed internally. Blocks with no decomposed - site run the plain frozen path — so a subset decomposition just leaves the rest - frozen. - - `sites` / `leading_axes` are static config.""" - - embed: Float[Array, "vocab d"] - stacked: LlamaLayer # the per-layer weights stacked on a leading layer axis (the scan - # `xs`), stored pre-stacked: a saved jit input, never re-stacked inside a forward. - n_layer: int = eqx.field(static=True) - norm: Float[Array, " d"] - lm_head: Float[Array, "vocab d"] - inv_freq: Float[Array, " hd2"] - sites: tuple[SiteSpec, ...] = eqx.field(static=True) - leading_axes: tuple[str, ...] = eqx.field(static=True) - eps: float = eqx.field(static=True) - scan_unroll: int = eqx.field(static=True, default=1) - """`lax.scan(unroll=)` factor over the block stack (`RuntimeConfig.scan_unroll`); 1 = - plain per-layer scan.""" - gather_fp8: bool = eqx.field(static=True, default=False) - """Quantized all-gather of the ÷fsdp compute V/U (`RuntimeConfig.gather_fp8`).""" - - @property - def site_names(self) -> tuple[str, ...]: - return tuple(s.name for s in self.sites) - - @property - def layers(self) -> list[LlamaLayer]: - """Per-layer view of `stacked` (slices the leading layer axis). For non-hot - consumers (attn-patterns recipe, equivalence harness); the forwards use `stacked`.""" - return [jax.tree.map(lambda a, idx=i: a[idx], self.stacked) for i in range(self.n_layer)] - - def shardings(self, mesh: "Mesh") -> "LlamaDecomposedModel": - """FSDP-on-`fsdp` the per-layer weights (`stacked.shardings` — `d` on `fsdp`, - head/intermediate replicated; the ~14 GB layer bulk shards `/fsdp`, gathered per layer - inside the scan, on NVLink). embed / lm_head / norm / inv_freq REPLICATE — the ~2 GB - embed+head is small and vocab-parallel logits/lookup aren't worth the complexity. (The - old all-replicate justification — "the target is small vs activations" — is stale: at - the full 32-layer model the replicated target + its backward/remat copies dominate the - step's peak, which is what this shards away.)""" - repl = NamedSharding(mesh, P()) - return eqx.tree_at( - lambda m: (m.embed, m.norm, m.lm_head, m.inv_freq, m.stacked), - self, - (repl, repl, repl, repl, self.stacked.shardings(mesh)), - ) - - @staticmethod - def recon_loss_fn(masked_output: Array, clean_output: Array) -> Array: - return kl_per_position(masked_output, clean_output) - - def embed_tokens(self, tokens: Int[Array, "b t"]) -> Float[Array, "b t d"]: - return self.embed[tokens] - - def clean_output(self, inputs: Int[Array, "b t"]) -> Array: - """The all-frozen forward — the recon target (SPEC S3). A `lax.scan` over the block - stack so XLA compiles one block body instead of unrolling all 32 layers (the compile - fix for the full model; the scan reassociates float ops vs an unrolled loop, within - fp32 tolerance).""" - - def block(x: Array, layer: LlamaLayer) -> tuple[Array, None]: - x = x + layer.attn(rms_norm(x, layer.ln1, self.eps), self.inv_freq) - x = x + _clean_mlp_out(layer, rms_norm(x, layer.ln2, self.eps)) - return x, None - - x, _ = jax.lax.scan(block, self.embed_tokens(inputs), self.stacked) - x = rms_norm(x, self.norm, self.eps) - return x @ self.lm_head.T - - def read_activations( - self, inputs: Int[Array, "b t"], wanted: tuple[str, ...] - ) -> dict[str, Array]: - """Frozen-path activation accessor (CI input side, SPEC S4; harvest's per-site - matrix inputs). - - `wanted` keys are either `resid.{layer}` (residual stream ENTERING that block — the - chunkwise CI fn's `input_names`) or a decomposed SITE NAME (the activation entering - that site's weight on the frozen path: `q/k/v_proj` ← post-LN1 residual, `o_proj` ← - the attention output, `gate/up_proj` ← post-LN2 residual, `down_proj` ← - `silu(gate)·up`). The residual is threaded identically to `clean_output`; the - per-site intermediates come from the same RMSNorm/attn/MLP math. Stops once the last - requested key's block is fully covered (no wasted block compute past it).""" - wanted_set = frozenset(wanted) - last = max(_tap_layer(key) for key in wanted) - taps: dict[str, Array] = {} - x = self.embed_tokens(inputs) - for layer in range(self.n_layer): - block = jax.tree.map(lambda a, li=layer: a[li], self.stacked) - if f"resid.{layer}" in wanted_set: - taps[f"resid.{layer}"] = x - attn = block.attn - h1 = rms_norm(x, block.ln1, self.eps) - attn_y = attn.core(h1 @ attn.wq.T, h1 @ attn.wk.T, h1 @ attn.wv.T, self.inv_freq) - post_attn = x + attn_y @ attn.wo.T - mlp_in = rms_norm(post_attn, block.ln2, self.eps) - down_in = jax.nn.silu(mlp_in @ block.Wg.T) * (mlp_in @ block.Wu.T) - for kind, site_input in ( - ("q", h1), ("k", h1), ("v", h1), ("o", attn_y), - ("gate", mlp_in), ("up", mlp_in), ("down", down_in), - ): # fmt: skip - name = site_name(layer, kind) - if name in wanted_set: - taps[name] = site_input - x = post_attn + down_in @ block.Wd.T - if layer == last: - break - assert set(taps) == wanted_set, (sorted(taps), sorted(wanted)) - return taps - - def _run_masked_forward( - self, - prepared: dict[str, dict[str, Array]], - inputs: Int[Array, "b t"], - masks: dict[str, Array], - delta_masks: dict[str, Array], - routes: dict[str, Array] | None, - live: tuple[str, ...], - has_delta: bool, - remat: bool, - collect: dict[str, Array] | None, - stochastic: tuple[dict[str, Array], Array] | None = None, - collect_activations: dict[str, Array] | None = None, - ) -> Array: - """The masked decomposed forward shared by `masked_output` / `masked_output_stochastic` - / `masked_site_outputs` / `masked_component_activations` (SPEC §1.3, S2). `live_set` is - static at trace, so the forward runs as `[frozen prefix] → [live block] → [frozen suffix]` - static sub-scans (no per-site `lax.cond`): only the live block carries + gathers V/U. - `live`/`has_delta` are static; a non-None `collect` gathers per-live-site decomposed - OUTPUTS (`(x@V)*m@U + …`, SPEC S31), and a non-None `collect_activations` gathers - per-live-site component ACTIVATIONS `x@V` (`[*leading, C]`, mask-independent — the - pre-mask coefficient the arithmetic CI-grid eval visualizes). Assumes layer-aligned, - contiguous chunks (asserted below). - - `prepared` is the shared, stacked + ÷fsdp-reconstructed per-kind `(V, U)` from - `prepare_compute_weights` (built ONCE per step) — this fn only ATTACHES the per-forward - masks, so the ÷N→÷fsdp cross-node gather is not re-run here (SPEC unchanged; numerics - identical — the reconstruction is mask-independent and the same for every forward).""" - live_set = frozenset(live) - resid = self.embed_tokens(inputs) - leading = resid.shape[:-1] - if stochastic is not None: - ci_stacked, draw_key = stochastic - per_kind = _attach_per_kind_stochastic( - prepared, self.n_layer, leading, ci_stacked, draw_key, routes, live_set - ) - else: - per_kind = _attach_per_kind_masks( - prepared, self.n_layer, leading, masks, delta_masks, routes, live_set, has_delta - ) - decomposed_kinds = frozenset(per_kind) - want_collect = collect is not None - want_collect_acts = collect_activations is not None - - # STATIC liveness. `live_set` is known at trace, so the live/frozen choice per site needs - # NO runtime `lax.cond` — and removing it lets XLA pack + prefetch the V/U gathers (the - # cond was a scheduling/packing barrier). We assume LAYER-ALIGNED, CONTIGUOUS chunks - # (every production plan: `into_groups` with sites_per_chunk % n_decomposed_kinds == 0, - # and `one_chunk`); both are asserted below. The forward is then - # [frozen prefix] → [live block] → [frozen suffix], each a static sub-scan; only the live - # block carries V/U and gathers them. - def layer_is_live(layer: int) -> bool: - flags = {site_name(layer, kind) in live_set for kind in decomposed_kinds} - assert len(flags) == 1, ( - f"layer {layer} is partially live ({flags}); the segmented masked forward assumes " - f"layer-aligned chunks (sites_per_chunk % {len(decomposed_kinds)} == 0)" - ) - return flags.pop() - - live_layers = [layer for layer in range(self.n_layer) if layer_is_live(layer)] - if live_layers: - first_live, last_live = live_layers[0], live_layers[-1] + 1 - assert live_layers == list(range(first_live, last_live)), ( - f"live layers must be contiguous, got {live_layers}" - ) - else: - first_live = last_live = 0 - - def decomp_site(x_in: Array, W: Array, e: dict[str, Array]) -> Array: - v, u = e["V"], e["U"] - if "V_scale" in e: # fp8 QAG: gather the fp8 ÷fsdp weight to full d (½ bytes on the - # wire), THEN dequant to bf16 — the barrier keeps the convert after the gather so - # the collective moves fp8, not bf16. - v = dequantize_fp8( - jax.lax.optimization_barrier( - jax.lax.with_sharding_constraint(v, P(None, "tp")) - ), - e["V_scale"], - ) - u = dequantize_fp8( - jax.lax.optimization_barrier( - jax.lax.with_sharding_constraint(u, P("tp", None)) - ), - e["U_scale"], - ) - if "ci" in e: # stochastic recompute: draw source from the per-layer key and build the - # mask INLINE (recomputed in the backward, not held — the shared `ci` stack + tiny - # key replace the per-forward mask stack). - ci = e["ci"] - source = jax.random.uniform(e["src_key"], ci.shape, dtype=ci.dtype) - mask = ci + (1.0 - ci) * source - delta = ( - jax.random.uniform(e["delta_key"], ci.shape[:-1], dtype=ci.dtype) - if has_delta - else None - ) - return site_out(x_in, v, u, W, mask, delta, e.get("route")) - return site_out(x_in, v, u, W, e["mask"], e.get("delta"), e.get("route")) - - def masked_site( - x_in: Array, kind: str, W: Array, pk: dict[str, dict[str, Array]] - ) -> tuple[Array, Array | None, Array | None]: - # LIVE block only: every decomposed kind decomps (static — no cond); a kind absent - # from the decomposition stays frozen. - if kind not in decomposed_kinds: - return x_in @ W.T, None, None - e = pk[kind] - out = decomp_site(x_in, W, e) - act = ( - (x_in @ e["V"]) if want_collect_acts else None - ) # pre-mask coeff; see masked_component_activations - return out, (out if want_collect else None), act - - def live_block( - x: Array, layer_in: tuple[LlamaLayer, dict[str, dict[str, Array]]] - ) -> tuple[Array, tuple[dict[str, Array] | None, dict[str, Array] | None]]: - sl, pk = layer_in - attn = sl.attn - h1 = rms_norm(x, sl.ln1, self.eps) - q, qc, qa = masked_site(h1, "q", attn.wq, pk) - k, kc, ka = masked_site(h1, "k", attn.wk, pk) - v, vc, va = masked_site(h1, "v", attn.wv, pk) - attn_y = attn.core(q, k, v, self.inv_freq) - o, oc, oa = masked_site(attn_y, "o", attn.wo, pk) - post_attn = x + o - h2 = rms_norm(post_attn, sl.ln2, self.eps) - g, gc, ga = masked_site(h2, "gate", sl.Wg, pk) - u, uc, ua = masked_site(h2, "up", sl.Wu, pk) - d, dc, da = masked_site(jax.nn.silu(g) * u, "down", sl.Wd, pk) - x = post_attn + d - kinds = ("q", "k", "v", "o", "gate", "up", "down") - collected = ( - { - n: c - for n, c in zip(kinds, (qc, kc, vc, oc, gc, uc, dc), strict=True) - if c is not None - } - if want_collect - else None - ) - collected_acts = ( - { - n: a - for n, a in zip(kinds, (qa, ka, va, oa, ga, ua, da), strict=True) - if a is not None - } - if want_collect_acts - else None - ) - return x, (collected, collected_acts) - - def frozen_block(x: Array, sl: LlamaLayer) -> tuple[Array, None]: - # Bit-identical to a frozen `masked_site` branch (`x @ Wᵀ` per site), shared with - # `clean_output`. Carries NO V/U → a frozen segment gathers nothing. - x = x + sl.attn(rms_norm(x, sl.ln1, self.eps), self.inv_freq) - x = x + _clean_mlp_out(sl, rms_norm(x, sl.ln2, self.eps)) - return x, None - - # Per-LAYER checkpoint of the scan BODY in BOTH modes — `remat` controls ONLY whether the - # layer ACTIVATIONS are recomputed; it NEVER controls the ÷fsdp→full V/U gather. That - # gather is a NON-dot collective, so it is never a saved residual under either policy — it - # re-gathers in the backward, transient one layer at a time (without checkpoint XLA keeps - # every layer's full gathered V/U live across the scan → OOM). - # remat=True → nothing_saveable: recompute activations AND the gather (min memory). - # remat=False → dots_saveable: SAVE the activation matmuls (no batch dims here → they - # qualify) but still recompute the gather + cheap elementwise. Pure recompute either - # way; zero numerics change. - # `scan_unroll` (native `lax.scan(unroll=k)`) emits k iterations straight-line so XLA can - # prefetch gather(L+1) under matmul(L) — the overlap a 1-layer while-body denies. - policy = ( - jax.checkpoint_policies.nothing_saveable - if remat - else jax.checkpoint_policies.dots_saveable - ) - - def run_scan(body: Any, carry: Array, xs: Any) -> tuple[Array, Any]: - return jax.lax.scan( - jax.checkpoint(body, policy=policy), carry, xs, unroll=self.scan_unroll - ) - - def slice_layers(lo: int, hi: int) -> LlamaLayer: - return jax.tree.map(lambda a: a[lo:hi], self.stacked) - - x = resid - ys: tuple[dict[str, Array] | None, dict[str, Array] | None] | None = None - if first_live > 0: - x, _ = run_scan(frozen_block, x, slice_layers(0, first_live)) - if last_live > first_live: - pk_live = { - kind: {k: v[first_live:last_live] for k, v in e.items()} - for kind, e in per_kind.items() - } - x, ys = run_scan(live_block, x, (slice_layers(first_live, last_live), pk_live)) - if last_live < self.n_layer: - x, _ = run_scan(frozen_block, x, slice_layers(last_live, self.n_layer)) - - x = rms_norm(x, self.norm, self.eps) - logits = x @ self.lm_head.T - if collect is not None or collect_activations is not None: - assert ys is not None # requested -> the live block emitted the per-kind stacks - ys_out, ys_acts = ys - for sink, stacked in ((collect, ys_out), (collect_activations, ys_acts)): - if sink is None: - continue - assert stacked is not None # sink requested -> the live block emitted this stack - for site in live: - layer, kind = parse_site_name(site) - sink[site] = stacked[kind][layer - first_live] - return logits - - def prepare_compute_weights(self, vu: DecompVU) -> dict[str, dict[str, Array]]: - """Build the shared per-kind compute weights ONCE per step (SPEC unchanged): stack the - per-site V/U into the layer-stacked `[n_layer, …]` form and run the ÷N→÷fsdp cross-node - reconstruction + bf16 cast. The result is mask-independent and identical for every - forward in the step, so the engine builds it once and threads it into all - `masked_output` / `masked_site_outputs` calls — the cross-node gather then runs ONCE per - step (ENTRY) instead of once per forward (the per-forward re-gather was ~10 co-resident - copies of the ÷fsdp stack at peak).""" - return _reconstruct_compute_weights(_stack_per_kind_vu(vu, self.n_layer), self.gather_fp8) - - def masked_output( - self, - prepared: dict[str, dict[str, Array]], - inputs: Int[Array, "b t"], - masks: dict[str, Array], - delta_masks: dict[str, Array], - routes: dict[str, Array] | None, - live: tuple[str, ...], - has_delta: bool, - *, - remat: bool, - ) -> Array: - return self._run_masked_forward( - prepared, inputs, masks, delta_masks, routes, live, has_delta, remat, None, None - ) - - def stack_ci(self, ci_lower: dict[str, Array]) -> dict[str, Array]: - """Per-kind `[n_layer, *leading, C]` stack of the CI envelope, built ONCE per step and - shared across all stochastic recon forwards (`masked_output_stochastic`). The - StochasticReconCapable capability (SPEC unchanged — pure recompute restructuring).""" - return _stack_ci_per_kind(ci_lower, self.n_layer) - - def masked_output_stochastic( - self, - prepared: dict[str, dict[str, Array]], - inputs: Int[Array, "b t"], - ci_stacked: dict[str, Array], - draw_key: Array, - routes: dict[str, Array] | None, - live: tuple[str, ...], - has_delta: bool, - *, - remat: bool, - ) -> Array: - """Stochastic recon forward that RECOMPUTES masks in-block (memory win): the shared - `ci_stacked` + per-layer keys from `draw_key` replace the per-forward mask stack; each - live site draws `source = uniform(key)` and forms `mask = ci + (1−ci)·source` inside the - checkpointed block (faithful by checkpoint determinism). Same forward semantics as - `masked_output` with stochastic sources — only the masks' liverange changes.""" - return self._run_masked_forward( - prepared, inputs, {}, {}, routes, live, has_delta, remat, None, (ci_stacked, draw_key) - ) - - def masked_site_outputs( - self, - prepared: dict[str, dict[str, Array]], - inputs: Int[Array, "b t"], - masks: dict[str, Array], - delta_masks: dict[str, Array], - routes: dict[str, Array] | None, - live: tuple[str, ...], - has_delta: bool, - ) -> dict[str, Array]: - """Per-`live`-site decomposed output of the masked forward (SPEC S31). Runs the - exact `masked_output` forward, discards the logits, returns the collected outputs.""" - collect: dict[str, Array] = {} - self._run_masked_forward( - prepared, inputs, masks, delta_masks, routes, live, has_delta, False, collect, None - ) - assert set(collect) == set(live), (sorted(collect), sorted(live)) - return collect - - def masked_component_activations( - self, - prepared: dict[str, dict[str, Array]], - inputs: Int[Array, "b t"], - masks: dict[str, Array], - delta_masks: dict[str, Array], - routes: dict[str, Array] | None, - live: tuple[str, ...], - has_delta: bool, - ) -> dict[str, Array]: - """Per-`live`-site component activation `x@V` (`[*leading, C]`, the coefficient - BEFORE the per-component `*mask`) from the masked forward. A site's OWN mask does not - enter its `x@V`, but the masks passed still shape a downstream site's input (e.g. - `down`'s input is the masked `silu(gate)*up`) — so pass the masks the visualization - wants the forward to run under. For the arithmetic CI-grid eval's activation heatmaps; - LM-only, off the recon path.""" - collect_activations: dict[str, Array] = {} - self._run_masked_forward( - prepared, inputs, masks, delta_masks, routes, live, has_delta, False, None, - collect_activations=collect_activations, - ) # fmt: skip - assert set(collect_activations) == set(live), (sorted(collect_activations), sorted(live)) - return collect_activations - - def weight_deltas(self, vu: DecompVU) -> dict[str, Array]: - """fp32 `W − V@U` per site from fp32 masters (SPEC N2; faithfulness input).""" - out: dict[str, Array] = {} - for spec in self.sites: - layer, kind = parse_site_name(spec.name) - W = _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind) - V, U = vu.site(spec.name) - out[spec.name] = ( - W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T - ) - return out - - -# ----------------------------- HF weight loading ----------------------------- - - -def hf_snapshot_dir(model_name: str) -> Path: - """Newest local snapshot of `model_name`. `HF_HUB_CACHE` overrides; otherwise on - cluster (`DATA_MOUNT` set) the shared world-readable cache is the source — a home - `~/.cache` hub is silently mutable, and a wiped entry strands running jobs that - reload weights on requeue.""" - import os - - default_cache = ( - f"{os.environ['DATA_MOUNT']}/artifacts/hf_cache/hub" - if "DATA_MOUNT" in os.environ - else str(Path.home() / ".cache/huggingface/hub") - ) - cache = Path(os.environ.get("HF_HUB_CACHE", default_cache)) - repo = "models--" + model_name.replace("/", "--") - snaps = sorted((cache / repo / "snapshots").iterdir()) - assert snaps, f"no snapshot for {model_name} under {cache}" - return snaps[-1] - - -class _HFWeights: - """Lazy keyed access to the sharded safetensors of an HF Llama checkpoint.""" - - def __init__(self, snapshot: Path): - index = json.loads((snapshot / "model.safetensors.index.json").read_text()) - self._key_to_file = index["weight_map"] - self._snapshot = snapshot - self._open: dict[str, Any] = {} - - def get(self, key: str) -> Array: - fname = self._key_to_file[key] - if fname not in self._open: - self._open[fname] = safe_open(str(self._snapshot / fname), framework="numpy") - return jnp.asarray(np.array(self._open[fname].get_tensor(key)), dtype=DT) - - -def _load_attn(w: _HFWeights, i: int, cfg: LlamaConfig) -> FrozenAttn: +def _load_attn(w: HFWeights, i: int, cfg: LlamaConfig) -> FrozenAttn: pre = "model.layers" return FrozenAttn( wq=w.get(f"{pre}.{i}.self_attn.q_proj.weight"), @@ -987,74 +46,21 @@ def _load_attn(w: _HFWeights, i: int, cfg: LlamaConfig) -> FrozenAttn: ) -def _load_blocks(w: "_HFWeights", cfg: LlamaConfig) -> list[LlamaLayer]: - pre = "model.layers" - return [ - LlamaLayer( - ln1=w.get(f"{pre}.{i}.input_layernorm.weight"), - ln2=w.get(f"{pre}.{i}.post_attention_layernorm.weight"), - attn=_load_attn(w, i, cfg), - Wg=w.get(f"{pre}.{i}.mlp.gate_proj.weight"), - Wu=w.get(f"{pre}.{i}.mlp.up_proj.weight"), - Wd=w.get(f"{pre}.{i}.mlp.down_proj.weight"), - ) - for i in range(cfg.n_layer) - ] - - -def build_decomposed_lm( - embed: Array, - layers: list[LlamaLayer], - norm: Array, - lm_head: Array, - inv_freq: Array, - cfg: LlamaConfig, - sites: tuple[SiteSpec, ...], - scan_unroll: int = 1, - gather_fp8: bool = False, -) -> LlamaDecomposedModel: - """Assemble a `LlamaDecomposedModel` from the frozen full-model arrays + decomposition - config. `sites` must be canonical-ordered with dims matching `cfg`. `scan_unroll` / - `gather_fp8` are the `RuntimeConfig` compute knobs (1 / off = the default forward).""" - site_cs = tuple(SiteC(s.name, s.C) for s in sites) - assert sites == llama_site_specs(cfg, canonical_site_cs(site_cs)), ( - f"sites are not the canonical specs for this config: {sites}" - ) - return LlamaDecomposedModel( - embed=embed, - stacked=_stack_layers(layers), - n_layer=len(layers), - norm=norm, - lm_head=lm_head, - inv_freq=inv_freq, - sites=sites, - leading_axes=("sequence",), - eps=cfg.rms_norm_eps, - scan_unroll=scan_unroll, - gather_fp8=gather_fp8, - ) - - -def load_decomposed_lm_from_hf( +def load_decomposed_llama_from_hf( model_name: str, cfg: LlamaConfig, sites: tuple[SiteSpec, ...], scan_unroll: int = 1, gather_fp8: bool = False, -) -> LlamaDecomposedModel: - """Load the Llama-8B `DecomposedModel`: the full frozen model (embedding, all blocks, - final norm, lm_head) as fields plus the static decomposition config (`sites`). Blocks - without a decomposed site run the plain frozen path. `scan_unroll` / `gather_fp8` are the - `RuntimeConfig` compute knobs.""" - w = _HFWeights(hf_snapshot_dir(model_name)) - return build_decomposed_lm( - embed=w.get("model.embed_tokens.weight"), - layers=_load_blocks(w, cfg), - norm=w.get("model.norm.weight"), - lm_head=w.get("lm_head.weight"), +) -> GLUDecomposedModel: + """The Llama family HF load: plain attention + llama3-rescaled RoPE frequencies.""" + return load_decomposed_glu_from_hf( + model_name, + cfg, + sites, + load_attn=lambda w, i: _load_attn(w, i, cfg), + weights_dtype=jnp.bfloat16, # the family is bf16-only (TargetConfig.supported_weights_dtypes) inv_freq=llama3_inv_freq(cfg), - cfg=cfg, - sites=sites, scan_unroll=scan_unroll, gather_fp8=gather_fp8, ) diff --git a/param_decomp/targets/llama_simple_mlp.py b/param_decomp/targets/llama_simple_mlp.py index dc71e6e63..7c177613f 100644 --- a/param_decomp/targets/llama_simple_mlp.py +++ b/param_decomp/targets/llama_simple_mlp.py @@ -43,7 +43,7 @@ from param_decomp.components import DecompVU, SiteC, SiteSpec, site_out from param_decomp.lm import run_stochastic_masked_output from param_decomp.losses import kl_per_position -from param_decomp.targets.llama8b import FrozenAttn +from param_decomp.targets.glu_transformer import FrozenAttn from vendored_jax.llama import rms_norm KIND_ORDER = ("q_proj", "k_proj", "v_proj", "o_proj", "c_fc", "down_proj") @@ -309,6 +309,17 @@ def shardings(self, mesh: Mesh) -> "SimpleMLPDecomposedModel": repl = NamedSharding(mesh, P()) return jax.tree.map(lambda _a: repl, self) + def attn_pattern( + self, + q_site: str, + q_flat: Float[Array, "b t qd"], + k_flat: Float[Array, "b t kvd"], + ) -> Float[Array, "b h t t"]: + """The target-owned attn-patterns eval recipe (`attn_patterns_eval`'s + `AttnPatternModel` protocol): that layer's own attention module's pattern.""" + layer = parse_site_name(q_site)[0] + return self.layers[layer].attn.pattern(q_flat, k_flat, self.inv_freq) + @staticmethod def recon_loss_fn(masked_output: Array, clean_output: Array) -> Array: return kl_per_position(masked_output, clean_output) diff --git a/param_decomp/targets/qwen3_8b.py b/param_decomp/targets/qwen3_8b.py new file mode 100644 index 000000000..a86cd74e5 --- /dev/null +++ b/param_decomp/targets/qwen3_8b.py @@ -0,0 +1,107 @@ +"""The Qwen3-8B-Base family target over the shared GLU-transformer machinery +(`targets/glu_transformer.py`). Qwen3's one structural delta vs Llama is QK-norm: +per-head RMSNorm on q/k between the projection and RoPE (HF `Qwen3Attention.q_norm` / +`k_norm`) — `Qwen3FrozenAttn` carries the norm weights as REQUIRED fields (never an +optional flag) and applies them in the `_prep_qk` hook. Decomposition semantics are +unchanged: a masked q/k site output feeds q_norm → RoPE → SDPA. + +JAX↔HF parity is pinned directly by `param_decomp/tests/qwen3_hf_parity/`.""" + +from typing import override + +import equinox as eqx +import jax.numpy as jnp +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P +from jaxtyping import Array, Float + +from param_decomp.components import SiteSpec +from param_decomp.targets.glu_transformer import ( + FrozenAttn, + GLUConfig, + GLUDecomposedModel, + HFWeights, + default_inv_freq, + load_decomposed_glu_from_hf, +) +from vendored_jax.llama import rms_norm + + +def qwen3_8b_config() -> GLUConfig: + """Qwen3-8B(-Base) per HF `Qwen/Qwen3-8B-Base` config.json. HF's explicit `head_dim` + (128) coincides with `n_embd // n_head` for 8B, so the derived `GLUConfig.head_dim` + is exact (NOT true of every Qwen3 size — check before adding one).""" + return GLUConfig( + vocab_size=151936, + n_layer=36, + n_head=32, + n_kv_head=8, + n_embd=4096, + n_intermediate=12288, + rope_theta=1000000.0, + rms_norm_eps=1e-6, + max_position_embeddings=32768, + ) + + +class Qwen3FrozenAttn(FrozenAttn): + """GQA attention + Qwen3 QK-norm (per-head RMSNorm over head_dim, before RoPE).""" + + q_norm: Float[Array, " hd"] + k_norm: Float[Array, " hd"] + eps: float = eqx.field(static=True) + """RMSNorm eps for the QK-norm only (block norms use the model-level eps).""" + + def __check_init__(self): + assert self.eps > 0.0, "QK-norm needs a real eps" + + @override + def _prep_qk( + self, q: Float[Array, "b t h hd"], k: Float[Array, "b t kvh hd"] + ) -> tuple[Array, Array]: + return rms_norm(q, self.q_norm, self.eps), rms_norm(k, self.k_norm, self.eps) + + @override + def shardings(self, mesh: Mesh) -> "Qwen3FrozenAttn": + """The shared projection layout plus replicated norm vectors.""" + repl = NamedSharding(mesh, P()) + placed = super().shardings(mesh) + assert isinstance(placed, Qwen3FrozenAttn) + return eqx.tree_at(lambda a: (a.q_norm, a.k_norm), placed, (repl, repl)) + + +def _load_attn(w: HFWeights, i: int, cfg: GLUConfig) -> Qwen3FrozenAttn: + pre = "model.layers" + return Qwen3FrozenAttn( + wq=w.get(f"{pre}.{i}.self_attn.q_proj.weight"), + wk=w.get(f"{pre}.{i}.self_attn.k_proj.weight"), + wv=w.get(f"{pre}.{i}.self_attn.v_proj.weight"), + wo=w.get(f"{pre}.{i}.self_attn.o_proj.weight"), + n_head=cfg.n_head, + n_kv_head=cfg.n_kv_head, + head_dim=cfg.head_dim, + n_rep=cfg.n_rep, + q_norm=w.get(f"{pre}.{i}.self_attn.q_norm.weight"), + k_norm=w.get(f"{pre}.{i}.self_attn.k_norm.weight"), + eps=cfg.rms_norm_eps, + ) + + +def load_decomposed_qwen3_from_hf( + model_name: str, + cfg: GLUConfig, + sites: tuple[SiteSpec, ...], + scan_unroll: int = 1, + gather_fp8: bool = False, +) -> GLUDecomposedModel: + """The Qwen3 family HF load: QK-norm attention + plain RoPE frequencies.""" + return load_decomposed_glu_from_hf( + model_name, + cfg, + sites, + load_attn=lambda w, i: _load_attn(w, i, cfg), + weights_dtype=jnp.bfloat16, # the family is bf16-only (TargetConfig.supported_weights_dtypes) + inv_freq=default_inv_freq(cfg.head_dim, cfg.rope_theta), + scan_unroll=scan_unroll, + gather_fp8=gather_fp8, + ) diff --git a/param_decomp/targets/target_aliases.py b/param_decomp/targets/target_aliases.py deleted file mode 100644 index 1af3ecdda..000000000 --- a/param_decomp/targets/target_aliases.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Decomposed-model union for the LM targets `run.py::main` and -`load_run.py::build_target` dispatch over. - -The toy targets (TMS, ResidMLP) live in the lab and are NOT members of this union — the -generic engine (`run_decomposition_training`) takes the model as `DecomposedModel`, so the -core never names a toy.""" - -from param_decomp.targets.llama8b import LlamaDecomposedModel -from param_decomp.targets.llama_simple_mlp import SimpleMLPDecomposedModel - -AnyDecomposedModel = LlamaDecomposedModel | SimpleMLPDecomposedModel diff --git a/param_decomp/tests/equivalence/jax_equivalence.py b/param_decomp/tests/equivalence/jax_equivalence.py index cada369c6..584315c31 100644 --- a/param_decomp/tests/equivalence/jax_equivalence.py +++ b/param_decomp/tests/equivalence/jax_equivalence.py @@ -34,14 +34,14 @@ importance_minimality_terms, kl_per_position, ) -from param_decomp.targets.llama8b import ( # noqa: E402 +from param_decomp.targets.glu_transformer import ( # noqa: E402 MLP_KINDS, FrozenAttn, - LlamaDecomposedModel, - LlamaLayer, + GLUDecomposedModel, + GLULayer, _clean_mlp_out, # noqa: E402 (reference suffix forward in the chunk-plan gate check) build_decomposed_lm, - llama_site_specs, + glu_site_specs, mlp_family_site_cs, site_name, ) @@ -81,7 +81,7 @@ def _build(f: dict[str, np.ndarray]): C = int(f[f"Vg_0"].shape[-1]) # noqa: F541 decomp_layers = [ - LlamaLayer( + GLULayer( ln1=a(f"ln1_{i}"), ln2=a(f"ln2_{i}"), attn=_zero_attn(d, di), @@ -92,7 +92,7 @@ def _build(f: dict[str, np.ndarray]): for i in range(n_layers) ] tail = [ - LlamaLayer( + GLULayer( ln1=a(f"tail_ln1_{j}"), ln2=a(f"tail_ln2_{j}"), attn=_zero_attn(d, di), @@ -133,7 +133,7 @@ def _build(f: dict[str, np.ndarray]): lm_head=a("lm_head"), inv_freq=inv_freq, cfg=cfg, - sites=llama_site_specs(cfg, mlp_family_site_cs(0, n_layers - 1, C)), + sites=glu_site_specs(cfg, mlp_family_site_cs(0, n_layers - 1, C)), ) return lm, vu, n_layers @@ -210,7 +210,7 @@ def per_site(prefix: str) -> dict[str, jnp.ndarray]: def _suffix_with_split_mlp( - tgt: LlamaDecomposedModel, + tgt: GLUDecomposedModel, vu: DecompVU, resid: jnp.ndarray, live_layer: int, diff --git a/param_decomp/tests/equivalence/test_equivalence.py b/param_decomp/tests/equivalence/test_equivalence.py index 51873da41..0e7557b89 100644 --- a/param_decomp/tests/equivalence/test_equivalence.py +++ b/param_decomp/tests/equivalence/test_equivalence.py @@ -185,7 +185,7 @@ def test_sc_source_broadcasts_over_batch_in_masked_forward() -> None: `delta_mask[..., None]` / mask broadcast (`components.site_out`) the way the PPGD path does, and pins the broadcast AXIS: transposing the source to `(1, B, C+1)` (B != T) must break the forward rather than silently re-interpret the time axis as batch.""" - from param_decomp.targets.llama8b import MLP_KINDS, site_name + from param_decomp.targets.glu_transformer import MLP_KINDS, site_name from param_decomp.tests.equivalence.jax_equivalence import FP, _build f = _load_fixtures() diff --git a/param_decomp/tests/qwen3_hf_parity/__init__.py b/param_decomp/tests/qwen3_hf_parity/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/param_decomp/tests/qwen3_hf_parity/gen_hf_fixtures.py b/param_decomp/tests/qwen3_hf_parity/gen_hf_fixtures.py new file mode 100644 index 000000000..ff2bc79fc --- /dev/null +++ b/param_decomp/tests/qwen3_hf_parity/gen_hf_fixtures.py @@ -0,0 +1,106 @@ +"""Generate the Qwen3 HF-parity goldens — a TORCH-ENV script (the repo venv is +torch-free; run it in a throwaway venv, like the `tests/equivalence` torch generators): + + uv venv /tmp/qwen3-golden --python 3.12 + VIRTUAL_ENV=/tmp/qwen3-golden uv pip install torch --index-url https://download.pytorch.org/whl/cpu + VIRTUAL_ENV=/tmp/qwen3-golden uv pip install transformers numpy + /tmp/qwen3-golden/bin/python param_decomp/tests/qwen3_hf_parity/gen_hf_fixtures.py # tiny + /tmp/qwen3-golden/bin/python param_decomp/tests/qwen3_hf_parity/gen_hf_fixtures.py --real # 8B + +Tiny golden (`qwen3_tiny_hf_fixtures.npz`): a seeded random `Qwen3ForCausalLM` at a +4-layer toy config, fp32, eager attention — the exact-architecture check (QK-norm, GQA, +RoPE base) `test_qwen3_hf_parity.py` compares against at fp32 tolerance. The full state +dict rides in the npz under `sd::`-prefixed keys. + +Real golden (`qwen3_8b_real_logits.npz`): `Qwen/Qwen3-8B-Base` bf16 from the HF cache, +a few fixed prompts, final-position fp32 logits — the weight-loading end-to-end check +(slow test). Regenerate only if the MATH changes; record the transformers version bump +in the commit message. +""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM + +HERE = Path(__file__).resolve().parent + +TINY_CONFIG = dict( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + rope_theta=1000000.0, + rms_norm_eps=1e-6, + max_position_embeddings=512, + tie_word_embeddings=False, + attention_bias=False, + use_cache=False, +) + +REAL_MODEL = "Qwen/Qwen3-8B-Base" +REAL_PROMPTS = ( + "The capital of France is Paris, and the capital of Germany is", + "In 1859, Charles Darwin published On the Origin of Species, which", + "def fibonacci(n):\n if n <= 1:\n return n\n return", + "The mitochondria is the powerhouse of the cell, and the nucleus stores the", +) +REAL_SEQ_LEN = 12 + + +def gen_tiny() -> None: + torch.manual_seed(0) + cfg = Qwen3Config(**TINY_CONFIG) + model = Qwen3ForCausalLM._from_config(cfg, attn_implementation="eager").eval() + tokens = torch.randint(0, cfg.vocab_size, (2, 16), generator=torch.Generator().manual_seed(1)) + with torch.no_grad(): + logits = model(tokens).logits + arrays = {f"sd::{k}": v.numpy() for k, v in model.state_dict().items()} + import transformers + + np.savez_compressed( + HERE / "qwen3_tiny_hf_fixtures.npz", + **arrays, + tokens=tokens.numpy(), + logits=logits.numpy(), + config_json=np.array(json.dumps(TINY_CONFIG)), + transformers_version=np.array(transformers.__version__), + ) + print(f"tiny golden: {len(arrays)} tensors, logits {tuple(logits.shape)}") + + +def gen_real() -> None: + tokenizer = AutoTokenizer.from_pretrained(REAL_MODEL) + ids = [] + for prompt in REAL_PROMPTS: + prompt_ids = tokenizer(prompt).input_ids + assert len(prompt_ids) >= REAL_SEQ_LEN, (prompt, len(prompt_ids)) + ids.append(prompt_ids[:REAL_SEQ_LEN]) + tokens = torch.tensor(ids) + model = Qwen3ForCausalLM.from_pretrained(REAL_MODEL, dtype=torch.bfloat16).eval() + with torch.no_grad(): + final_logits = model(tokens).logits[:, -1, :].float() + import transformers + + np.savez_compressed( + HERE / "qwen3_8b_real_logits.npz", + tokens=tokens.numpy(), + final_logits=final_logits.numpy(), + transformers_version=np.array(transformers.__version__), + ) + print(f"real golden: tokens {tuple(tokens.shape)}, final logits {tuple(final_logits.shape)}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--real", action="store_true", help="generate the Qwen3-8B-Base golden") + if parser.parse_args().real: + gen_real() + else: + gen_tiny() diff --git a/param_decomp/tests/qwen3_hf_parity/qwen3_8b_real_logits.npz b/param_decomp/tests/qwen3_hf_parity/qwen3_8b_real_logits.npz new file mode 100644 index 000000000..6e361205c Binary files /dev/null and b/param_decomp/tests/qwen3_hf_parity/qwen3_8b_real_logits.npz differ diff --git a/param_decomp/tests/qwen3_hf_parity/qwen3_tiny_hf_fixtures.npz b/param_decomp/tests/qwen3_hf_parity/qwen3_tiny_hf_fixtures.npz new file mode 100644 index 000000000..cd0ee23e6 Binary files /dev/null and b/param_decomp/tests/qwen3_hf_parity/qwen3_tiny_hf_fixtures.npz differ diff --git a/param_decomp/tests/qwen3_hf_parity/test_qwen3_hf_parity.py b/param_decomp/tests/qwen3_hf_parity/test_qwen3_hf_parity.py new file mode 100644 index 000000000..3c3ed90bf --- /dev/null +++ b/param_decomp/tests/qwen3_hf_parity/test_qwen3_hf_parity.py @@ -0,0 +1,128 @@ +"""Direct JAX-vs-HuggingFace parity for the Qwen3 target. + +The tiny test rebuilds `gen_hf_fixtures.py`'s seeded random `Qwen3ForCausalLM` as a +`GLUDecomposedModel` (Qwen3 family) from the golden's own state dict, fp32, and matches HF's +logits at fp32 tolerance — the exact-architecture check (QK-norm, GQA, plain RoPE at +theta 1e6). The slow test loads the REAL `Qwen/Qwen3-8B-Base` snapshot through the +production loader (`load_decomposed_lm_from_hf`, bf16) and matches HF's bf16 +final-position logits in distribution (KL + argmax) — the weight-loading check. +""" + +import json +from pathlib import Path + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from param_decomp.targets.glu_transformer import ( + GLUConfig, + GLUDecomposedModel, + GLULayer, + build_decomposed_lm, + default_inv_freq, + hf_snapshot_dir, +) +from param_decomp.targets.qwen3_8b import ( + Qwen3FrozenAttn, + load_decomposed_qwen3_from_hf, + qwen3_8b_config, +) + +HERE = Path(__file__).resolve().parent + + +def _tiny_cfg_from_golden(config_json: str) -> GLUConfig: + hf = json.loads(config_json) + assert hf["head_dim"] * hf["num_attention_heads"] == hf["hidden_size"], hf + assert not hf["tie_word_embeddings"] and not hf["attention_bias"], hf + return GLUConfig( + vocab_size=hf["vocab_size"], + n_layer=hf["num_hidden_layers"], + n_head=hf["num_attention_heads"], + n_kv_head=hf["num_key_value_heads"], + n_embd=hf["hidden_size"], + n_intermediate=hf["intermediate_size"], + rope_theta=hf["rope_theta"], + rms_norm_eps=hf["rms_norm_eps"], + max_position_embeddings=hf["max_position_embeddings"], + ) + + +def _build_from_hf_state(cfg: GLUConfig, sd: dict[str, np.ndarray]) -> GLUDecomposedModel: + """The Qwen3 family loader mirrored over an in-memory fp32 HF state dict (the + production loader reads sharded safetensors and casts bf16; here we stay fp32 so the + comparison isolates the MATH from rounding).""" + a = lambda key: jnp.asarray(sd[key], jnp.float32) # noqa: E731 + pre = "model.layers" + layers = [ + GLULayer( + ln1=a(f"{pre}.{i}.input_layernorm.weight"), + ln2=a(f"{pre}.{i}.post_attention_layernorm.weight"), + attn=Qwen3FrozenAttn( + wq=a(f"{pre}.{i}.self_attn.q_proj.weight"), + wk=a(f"{pre}.{i}.self_attn.k_proj.weight"), + wv=a(f"{pre}.{i}.self_attn.v_proj.weight"), + wo=a(f"{pre}.{i}.self_attn.o_proj.weight"), + n_head=cfg.n_head, + n_kv_head=cfg.n_kv_head, + head_dim=cfg.head_dim, + n_rep=cfg.n_rep, + q_norm=a(f"{pre}.{i}.self_attn.q_norm.weight"), + k_norm=a(f"{pre}.{i}.self_attn.k_norm.weight"), + eps=cfg.rms_norm_eps, + ), + Wg=a(f"{pre}.{i}.mlp.gate_proj.weight"), + Wu=a(f"{pre}.{i}.mlp.up_proj.weight"), + Wd=a(f"{pre}.{i}.mlp.down_proj.weight"), + ) + for i in range(cfg.n_layer) + ] + return build_decomposed_lm( + embed=a("model.embed_tokens.weight"), + layers=layers, + norm=a("model.norm.weight"), + lm_head=a("lm_head.weight"), + inv_freq=default_inv_freq(cfg.head_dim, cfg.rope_theta), + cfg=cfg, + sites=(), + ) + + +def test_tiny_random_qwen3_matches_hf(): + f = np.load(HERE / "qwen3_tiny_hf_fixtures.npz") + cfg = _tiny_cfg_from_golden(str(f["config_json"])) + sd = {k.removeprefix("sd::"): f[k] for k in f.files if k.startswith("sd::")} + lm = _build_from_hf_state(cfg, sd) + logits = np.asarray(lm.clean_output(jnp.asarray(f["tokens"]))) + np.testing.assert_allclose(logits, f["logits"], rtol=2e-4, atol=1e-5) + + +@pytest.mark.slow +def test_real_qwen3_8b_matches_hf(): + """The production HF loader against real-weights HF logits: bf16 both sides, so the + comparison is distributional — small KL at the final position and argmax agreement — + not elementwise.""" + try: + hf_snapshot_dir("Qwen/Qwen3-8B-Base") + except (AssertionError, FileNotFoundError, KeyError): + pytest.skip("no local Qwen/Qwen3-8B-Base snapshot") + f = np.load(HERE / "qwen3_8b_real_logits.npz") + lm = load_decomposed_qwen3_from_hf("Qwen/Qwen3-8B-Base", qwen3_8b_config(), ()) + logits = np.asarray(lm.clean_output(jnp.asarray(f["tokens"]))[:, -1, :].astype(jnp.float32)) + ref = f["final_logits"] + # tie-aware argmax: the golden is bf16, so distinct plausible tokens often carry the + # IDENTICAL quantized top logit (observed: a 3-way 32.25 tie) — require JAX's argmax + # to sit at the golden's top logit level, not to match one arbitrary tie-winner. + jax_pick_ref_logit = np.take_along_axis(ref, logits.argmax(-1, keepdims=True), -1)[:, 0] + assert (jax_pick_ref_logit >= ref.max(-1) - 0.25).all(), ( + logits.argmax(-1), + ref.argmax(-1), + jax_pick_ref_logit, + ref.max(-1), + ) + ref_logp = np.asarray(jax.nn.log_softmax(jnp.asarray(ref), axis=-1)) + jax_logp = np.asarray(jax.nn.log_softmax(jnp.asarray(logits), axis=-1)) + kl = (np.exp(ref_logp) * (ref_logp - jax_logp)).sum(-1) + assert (kl < 5e-3).all(), kl diff --git a/param_decomp/tests/stacked_parity/test_stacked_parity.py b/param_decomp/tests/stacked_parity/test_stacked_parity.py index 4d05ec4d6..fb11be621 100644 --- a/param_decomp/tests/stacked_parity/test_stacked_parity.py +++ b/param_decomp/tests/stacked_parity/test_stacked_parity.py @@ -47,11 +47,11 @@ from param_decomp.lm import DecomposedModel from param_decomp.recon import StochasticSources, build_loss_terms, subset_chunk_plan from param_decomp.schedule import ScheduleConfig -from param_decomp.targets.llama8b import ( +from param_decomp.targets.glu_transformer import ( FrozenAttn, - LlamaLayer, + GLULayer, build_decomposed_lm, - llama_site_specs, + glu_site_specs, mlp_family_site_cs, ) from param_decomp.tests.test_llama8b import _tiny_cfg @@ -95,7 +95,7 @@ def a(key: str) -> jnp.ndarray: return jnp.asarray(f[key]) layers = [ - LlamaLayer( + GLULayer( ln1=a(f"tgt::layers.{i}.ln1"), ln2=a(f"tgt::layers.{i}.ln2"), attn=FrozenAttn( @@ -114,7 +114,7 @@ def a(key: str) -> jnp.ndarray: ) for i in range(first, cfg.n_layer) ] - sites = llama_site_specs(cfg, mlp_family_site_cs(first, last, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(first, last, C)) lm = build_decomposed_lm( embed=jnp.zeros((cfg.vocab_size, cfg.n_embd), jnp.float32), layers=layers, norm=a("tgt::norm"), lm_head=a("tgt::lm_head"), diff --git a/param_decomp/tests/test_arithmetic_eval.py b/param_decomp/tests/test_arithmetic_eval.py index e3fc1a52f..874990e0b 100644 --- a/param_decomp/tests/test_arithmetic_eval.py +++ b/param_decomp/tests/test_arithmetic_eval.py @@ -25,7 +25,7 @@ ) from param_decomp.ci_fn import lower_leaky_hard_sigmoid from param_decomp.components import init_decomp_vu -from param_decomp.targets.llama8b import llama_site_specs, mlp_family_site_cs +from param_decomp.targets.glu_transformer import glu_site_specs, mlp_family_site_cs from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm from param_decomp.tests.test_slow_eval import _build_ci_fn from param_decomp.train import COMPUTE_DT, cast_floating @@ -39,7 +39,7 @@ def _tiny_setup(): cfg = _tiny_cfg() C = 8 - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) ci_fn = _build_ci_fn(lm, cfg.n_embd, jax.random.PRNGKey(2)) return cfg, lm, ci_fn, C diff --git a/param_decomp/tests/test_attn_patterns_eval.py b/param_decomp/tests/test_attn_patterns_eval.py index 0493b5fe7..2984648e6 100644 --- a/param_decomp/tests/test_attn_patterns_eval.py +++ b/param_decomp/tests/test_attn_patterns_eval.py @@ -1,7 +1,7 @@ """CPU tests for the in-loop attention-pattern recon eval metrics. -Pins the metric-local `attn_pattern_for` target dispatch (shape + causal/softmax sanity on -both LM targets), the all-false-routes clean target (KL=0 when masked==clean), and the +Pins the target-owned `attn_pattern` recipe (shape + causal/softmax sanity on both LM +targets), the all-false-routes clean target (KL=0 when masked==clean), and the host-side token-weighted accumulation (combined = Σ sum_kl / Σ n_distributions). """ @@ -14,7 +14,6 @@ from param_decomp.attn_patterns_eval import ( accumulate_attn_patterns, - attn_pattern_for, attn_patterns_log_entries, make_ci_attn_patterns_step, make_stochastic_attn_patterns_step, @@ -27,9 +26,7 @@ ) from param_decomp.components import SiteC, SiteSpec, init_decomp_vu from param_decomp.lm import DecomposedModel, run_stochastic_masked_output -from param_decomp.targets.llama8b import ( - llama_site_specs, -) +from param_decomp.targets.glu_transformer import glu_site_specs from param_decomp.targets.llama_simple_mlp import ( canonical_site_cs as simple_canonical, ) @@ -66,16 +63,15 @@ def _build_ci_fn(lm: DecomposedModel, n_embd: int, key: jax.Array) -> CIFn: return build_ci_fn(arch, lm.sites, key) -def test_attn_pattern_for_shape_and_causal_softmax_llama(): +def test_attn_pattern_shape_and_causal_softmax_llama(): cfg = _llama_cfg() - sites = llama_site_specs(cfg, (SiteC("layers.0.self_attn.q_proj", 4),)) + sites = glu_site_specs(cfg, (SiteC("layers.0.self_attn.q_proj", 4),)) lm = _llama_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) - pattern_fn = attn_pattern_for(lm) b, t = 2, 9 qd, kvd = cfg.n_head * cfg.head_dim, cfg.n_kv_head * cfg.head_dim q = jax.random.normal(jax.random.PRNGKey(1), (b, t, qd)) k = jax.random.normal(jax.random.PRNGKey(2), (b, t, kvd)) - pattern = np.asarray(pattern_fn(q, k)) + pattern = np.asarray(lm.attn_pattern("layers.0.self_attn.q_proj", q, k)) assert pattern.shape == (b, cfg.n_head, t, t) np.testing.assert_allclose(pattern.sum(-1), 1.0, rtol=1e-5, atol=1e-5) @@ -84,16 +80,15 @@ def test_attn_pattern_for_shape_and_causal_softmax_llama(): assert pattern.dtype == np.float32 -def test_attn_pattern_for_shape_and_causal_softmax_simple_mlp(): +def test_attn_pattern_shape_and_causal_softmax_simple_mlp(): cfg = _simple_cfg() sites = simple_site_specs(cfg, (SiteC("h.0.attn.q_proj", 4),)) lm = _simple_decomposed_model(cfg, sites, jax.random.PRNGKey(0)) - pattern_fn = attn_pattern_for(lm) b, t = 2, 7 qd, kvd = cfg.n_head * cfg.head_dim, cfg.n_kv_head * cfg.head_dim q = jax.random.normal(jax.random.PRNGKey(1), (b, t, qd)) k = jax.random.normal(jax.random.PRNGKey(2), (b, t, kvd)) - pattern = np.asarray(pattern_fn(q, k)) + pattern = np.asarray(lm.attn_pattern("h.0.attn.q_proj", q, k)) assert pattern.shape == (b, cfg.n_head, t, t) np.testing.assert_allclose(pattern.sum(-1), 1.0, rtol=1e-5, atol=1e-5) @@ -101,11 +96,6 @@ def test_attn_pattern_for_shape_and_causal_softmax_simple_mlp(): assert np.allclose(pattern[:, :, upper], 0.0) -def test_attn_pattern_for_refuses_non_attention_target(): - with pytest.raises(AssertionError, match="only applies to attention targets"): - attn_pattern_for(object()) - - def _llama_attn_setup(): cfg = _llama_cfg() site_cs = ( @@ -114,9 +104,9 @@ def _llama_attn_setup(): SiteC("layers.5.self_attn.q_proj", 8), SiteC("layers.5.self_attn.k_proj", 6), ) - from param_decomp.targets.llama8b import canonical_site_cs + from param_decomp.targets.glu_transformer import canonical_site_cs - sites = llama_site_specs(cfg, canonical_site_cs(site_cs)) + sites = glu_site_specs(cfg, canonical_site_cs(site_cs)) lm = _llama_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) components = init_decomp_vu(sites, jax.random.PRNGKey(1)) ci_fn = _build_ci_fn(lm, cfg.n_embd, jax.random.PRNGKey(2)) @@ -125,8 +115,7 @@ def _llama_attn_setup(): def test_ci_step_clean_equals_masked_when_ci_all_one_gives_finite_kl(): cfg, lm, components, ci_fn = _llama_attn_setup() - pattern_fn = attn_pattern_for(lm) - step = make_ci_attn_patterns_step(lm, pattern_fn) + step = make_ci_attn_patterns_step(lm) b, t = 2, 12 residual = jax.random.randint(jax.random.PRNGKey(4), (b, t), 0, cfg.vocab_size) @@ -142,8 +131,7 @@ def test_ci_step_clean_equals_masked_when_ci_all_one_gives_finite_kl(): def test_accumulate_is_token_weighted_and_combines(): cfg, lm, components, ci_fn = _llama_attn_setup() - pattern_fn = attn_pattern_for(lm) - step = make_ci_attn_patterns_step(lm, pattern_fn) + step = make_ci_attn_patterns_step(lm) res_a = jax.random.randint(jax.random.PRNGKey(4), (2, 10), 0, cfg.vocab_size) res_b = jax.random.randint(jax.random.PRNGKey(5), (2, 10), 0, cfg.vocab_size) @@ -173,9 +161,8 @@ def test_accumulate_is_token_weighted_and_combines(): def test_stochastic_step_runs_and_scales_n_by_draws(): cfg, lm, components, ci_fn = _llama_attn_setup() - pattern_fn = attn_pattern_for(lm) n_draws = 3 - step = make_stochastic_attn_patterns_step(lm, pattern_fn, n_draws) + step = make_stochastic_attn_patterns_step(lm, n_draws) b, t = 2, 8 residual = jax.random.randint(jax.random.PRNGKey(4), (b, t), 0, cfg.vocab_size) @@ -192,7 +179,7 @@ def test_simple_mlp_step_runs_end_to_end(): lm = _simple_decomposed_model(cfg, sites, jax.random.PRNGKey(0)) components = init_decomp_vu(sites, jax.random.PRNGKey(1)) ci_fn = _build_ci_fn(lm, cfg.n_embd, jax.random.PRNGKey(2)) - step = make_ci_attn_patterns_step(lm, attn_pattern_for(lm)) + step = make_ci_attn_patterns_step(lm) b, t = 2, 10 residual = jax.random.randint(jax.random.PRNGKey(4), (b, t), 0, cfg.vocab_size) @@ -293,8 +280,7 @@ def test_attn_patterns_steps_reject_positionless_target(): leading_axes=(), ) assert lm.leading_axes == () - dummy_pattern_fn = lambda q, k: q # noqa: E731 — never reached; assert fires first with pytest.raises(AssertionError, match="LM-only"): - make_ci_attn_patterns_step(lm, dummy_pattern_fn) + make_ci_attn_patterns_step(lm) with pytest.raises(AssertionError, match="LM-only"): - make_stochastic_attn_patterns_step(lm, dummy_pattern_fn, 1) + make_stochastic_attn_patterns_step(lm, 1) diff --git a/param_decomp/tests/test_checkpoint.py b/param_decomp/tests/test_checkpoint.py index ea4cddbf2..3e6a9a382 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -41,11 +41,11 @@ from param_decomp.recon import build_loss_terms from param_decomp.schedule import ScheduleConfig from param_decomp.sharding import hsdp_mesh -from param_decomp.targets.llama8b import ( - llama_site_specs, +from param_decomp.targets.glu_transformer import ( + glu_site_specs, mlp_family_site_cs, ) -from param_decomp.targets.llama8b_sharding import ( +from param_decomp.targets.glu_transformer_sharding import ( init_ci_fn_placed, init_decomp_vu_placed, init_sources_sharded, @@ -97,7 +97,7 @@ def _chunkwise_arch(lm: DecomposedModel, cfg: LlamaConfig) -> ChunkwiseTransform def _build(seed: int): cfg = _tiny_cfg() C, seq = 8, 16 - sites = llama_site_specs(cfg, mlp_family_site_cs(3, 4, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(3, 4, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu(sites, jax.random.PRNGKey(seed)) ci_fn = build_ci_fn(_chunkwise_arch(lm, cfg), lm.sites, jax.random.PRNGKey(seed + 1)) @@ -238,7 +238,7 @@ def _build_sharded(seed: int, mesh: Mesh): cfg = _tiny_cfg() n = mesh.devices.size C, seq = 8 * n, 16 - sites = llama_site_specs(cfg, mlp_family_site_cs(3, 4, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(3, 4, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu_placed(sites, jax.random.PRNGKey(seed), mesh) ci_fn = init_ci_fn_placed( diff --git a/param_decomp/tests/test_checkpoint_production_topology.py b/param_decomp/tests/test_checkpoint_production_topology.py index 85b3ef3a3..02d94dd37 100644 --- a/param_decomp/tests/test_checkpoint_production_topology.py +++ b/param_decomp/tests/test_checkpoint_production_topology.py @@ -8,7 +8,7 @@ topology: * the V/U + Adam states are C-SHARDED and the sources/moments REPLICATED over a - multi-device `dp` mesh (`llama8b_sharding.py`), exactly as `init_train_state` places + multi-device `dp` mesh (`glu_transformer_sharding.py`), exactly as `init_train_state` places them — so the test exercises the sharded save/restore path, not the all-on-one path; * MULTIPLE persistent terms (SPEC S23: one `adversaries` entry per term), so a per-term moment tree that got dropped would surface; @@ -50,8 +50,8 @@ from param_decomp.recon import build_loss_terms, persistent_configs from param_decomp.run import _ensure_global from param_decomp.schedule import ScheduleConfig -from param_decomp.targets.llama8b import llama_site_specs, mlp_family_site_cs -from param_decomp.targets.llama8b_sharding import ( +from param_decomp.targets.glu_transformer import glu_site_specs, mlp_family_site_cs +from param_decomp.targets.glu_transformer_sharding import ( hsdp_mesh, init_ci_fn_placed, init_decomp_vu_placed, @@ -89,7 +89,7 @@ def _build_sharded(seed: int): mesh = hsdp_mesh() cfg = _tiny_cfg() C, seq = 8, 16 - sites = llama_site_specs(cfg, mlp_family_site_cs(3, 4, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(3, 4, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) by_layer: dict[int, list[str]] = {} diff --git a/param_decomp/tests/test_config.py b/param_decomp/tests/test_config.py index 38a0519dd..ec9504dc8 100644 --- a/param_decomp/tests/test_config.py +++ b/param_decomp/tests/test_config.py @@ -21,7 +21,7 @@ build_loss_terms, persistent_configs, ) -from param_decomp.targets.llama8b import mlp_family_site_cs +from param_decomp.targets.glu_transformer import mlp_family_site_cs from param_decomp_lab.experiments.lm.config import ( LMExperimentConfig, assert_supported_weights_dtype, @@ -209,8 +209,8 @@ def with_ppgd_fields(**fields: object): def test_unsupported_model_family_refuses_and_supported_families_dispatch(): - """E23: only Llama-3.1-8B (`hf`/`hf_weights_in_vendored` - → `TargetConfig`) and `LlamaSimpleMLP` (`pretrained` → + """E23: only the `HF_MODEL_FAMILIES` models (`hf`/`hf_weights_in_vendored` + → `TargetConfig`; Llama-3.1-8B and Qwen3-8B-Base) and `LlamaSimpleMLP` (`pretrained` → `LlamaSimpleMLPTargetConfig`) convert; every other family is refused at convert time. The schema's `LMTargetSpec` discriminated union still validates a GPT-2 spec (it's a well-formed `kind`), so the refusal must come from `_resolve_target`'s @@ -243,6 +243,15 @@ def _converted_target(spec: dict[str, str]): ) assert isinstance(raw_hf_llama, TargetConfig) + raw_hf_qwen3 = _converted_target( + { + "kind": "hf", + "model_class": "transformers.Qwen3ForCausalLM", + "model_name": "Qwen/Qwen3-8B-Base", + } + ) + assert isinstance(raw_hf_qwen3, TargetConfig) + gpt2_hf = { "kind": "hf", "model_class": "transformers.GPT2LMHeadModel", diff --git a/param_decomp/tests/test_eval.py b/param_decomp/tests/test_eval.py index c28e45a98..624af8cc3 100644 --- a/param_decomp/tests/test_eval.py +++ b/param_decomp/tests/test_eval.py @@ -22,7 +22,7 @@ from param_decomp.components import SiteSpec from param_decomp.eval import make_eval_step, next_token_cross_entropy from param_decomp.lm import DecomposedModel, run_stochastic_masked_output -from param_decomp.targets.llama8b import llama_site_specs, mlp_family_site_cs +from param_decomp.targets.glu_transformer import glu_site_specs, mlp_family_site_cs from param_decomp.tests.test_llama8b import ( _tiny_cfg, _tiny_decomposed_lm, @@ -148,7 +148,7 @@ def test_next_token_cross_entropy_matches_manual(): def test_eval_step_keys_identities_and_determinism(): cfg = _tiny_cfg() C = 8 - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) from param_decomp.components import init_decomp_vu @@ -219,7 +219,7 @@ def test_eval_step_fresh_pgd_probe(): random source it starts from (ascent on a fixed objective), and be deterministic.""" cfg = _tiny_cfg() C = 8 - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 4, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 4, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) from param_decomp.components import init_decomp_vu @@ -281,7 +281,7 @@ def test_eval_step_fresh_pgd_probe_device_count_invariant(): n_dev = mesh.devices.size cfg = _tiny_cfg() - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 4, 8)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 4, 8)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) ci_fn = _build_ci_fn(lm, cfg.n_embd, jax.random.PRNGKey(2)) @@ -317,7 +317,7 @@ def test_eval_step_l0_groups_sum_member_sites(): """torch CI_L0 `groups` parity: a group's L0 is the SUM of its fnmatch-member sites' L0s; an unmatched pattern refuses at build time.""" cfg = _tiny_cfg() - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) from param_decomp.components import init_decomp_vu @@ -362,7 +362,7 @@ def test_eval_step_n_valid_rows_masks_pad_tail(): objective). The stochastic variants draw shape-dependent randomness, so they only agree in expectation and are excluded.""" cfg = _tiny_cfg() - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) from param_decomp.components import init_decomp_vu diff --git a/param_decomp/tests/test_fnmatch_site_order.py b/param_decomp/tests/test_fnmatch_site_order.py index 4955027a3..3cd37e1b3 100644 --- a/param_decomp/tests/test_fnmatch_site_order.py +++ b/param_decomp/tests/test_fnmatch_site_order.py @@ -26,7 +26,7 @@ import pytest from param_decomp.components import SiteC -from param_decomp.targets import llama8b, llama_simple_mlp +from param_decomp.targets import glu_transformer, llama_simple_mlp def _named_modules_order( @@ -72,8 +72,8 @@ def _llama8b_module_names(n_layer: int) -> tuple[str, ...]: "layers", "self_attn", "mlp", - tuple(f"{k}_proj" for k in llama8b.KIND_ORDER), - tuple(f"{k}_proj" for k in llama8b.ATTN_KINDS), + tuple(f"{k}_proj" for k in glu_transformer.KIND_ORDER), + tuple(f"{k}_proj" for k in glu_transformer.ATTN_KINDS), n_layer, ) @@ -99,7 +99,7 @@ def test_llama8b_single_layer_mlp_set_matches_torch(): SiteC("layers.18.mlp.up_proj", 24576), SiteC("layers.18.mlp.down_proj", 24576), ) - jax_sites = llama8b.canonical_site_cs(targets) + jax_sites = glu_transformer.canonical_site_cs(targets) torch_sites = _torch_resolve(targets, _llama8b_module_names(32)) assert set(jax_sites) == set(torch_sites) diff --git a/param_decomp/tests/test_fresh_pgd_cscope_dp_invariance.py b/param_decomp/tests/test_fresh_pgd_cscope_dp_invariance.py index 2e7b8c949..75b96f004 100644 --- a/param_decomp/tests/test_fresh_pgd_cscope_dp_invariance.py +++ b/param_decomp/tests/test_fresh_pgd_cscope_dp_invariance.py @@ -31,8 +31,8 @@ from param_decomp.components import init_decomp_vu from param_decomp.losses import kl_per_position from param_decomp.sharding import hsdp_mesh, shard_batch -from param_decomp.targets.llama8b import ( - llama_site_specs, +from param_decomp.targets.glu_transformer import ( + glu_site_specs, mlp_family_site_cs, ) from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm @@ -51,7 +51,7 @@ def _ascend_cscope_source( cfg = _tiny_cfg() first_layer = 3 C, seq, gbatch = 8, 16, 8 - sites = llama_site_specs(cfg, mlp_family_site_cs(first_layer, first_layer + 2, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(first_layer, first_layer + 2, C)) lm = _tiny_decomposed_lm(cfg, sites, random.PRNGKey(0)) components = jax.tree.map( lambda x: jax.lax.stop_gradient(x), init_decomp_vu(sites, random.PRNGKey(1)) diff --git a/param_decomp/tests/test_llama8b.py b/param_decomp/tests/test_llama8b.py index 00fb594b2..f8598c5e2 100644 --- a/param_decomp/tests/test_llama8b.py +++ b/param_decomp/tests/test_llama8b.py @@ -37,13 +37,13 @@ from param_decomp.lm import DecomposedModel from param_decomp.recon import build_loss_terms from param_decomp.schedule import ScheduleConfig -from param_decomp.targets.llama8b import ( +from param_decomp.targets.glu_transformer import ( FrozenAttn, - LlamaDecomposedModel, - LlamaLayer, + GLUDecomposedModel, + GLULayer, build_decomposed_lm, canonical_site_cs, - llama_site_specs, + glu_site_specs, mlp_family_site_cs, parse_site_name, site_name, @@ -72,8 +72,8 @@ def _tiny_cfg() -> LlamaConfig: def _tiny_decomposed_lm( cfg: LlamaConfig, sites: tuple[SiteSpec, ...], key: jax.Array -) -> LlamaDecomposedModel: - """A tiny random `LlamaDecomposedModel` (random embedding + full frozen layer stack +) -> GLUDecomposedModel: + """A tiny random `GLUDecomposedModel` (random embedding + full frozen layer stack plus the decomposition `sites`) — the CPU-test analog of `load_decomposed_lm_from_hf`.""" ks = iter(jax.random.split(key, 1024)) d, di = cfg.n_embd, cfg.n_intermediate @@ -89,9 +89,7 @@ def fattn(): ) # fmt: skip def layer(): - return LlamaLayer( - jnp.ones((d,)), jnp.ones((d,)), fattn(), n((di, d)), n((di, d)), n((d, di)) - ) + return GLULayer(jnp.ones((d,)), jnp.ones((d,)), fattn(), n((di, d)), n((di, d)), n((d, di))) return build_decomposed_lm( embed=n((cfg.vocab_size, d), 0.02), @@ -105,7 +103,7 @@ def layer(): def _mlp_sites(cfg: LlamaConfig, first: int, last: int, C: int) -> tuple[SiteSpec, ...]: - return llama_site_specs(cfg, mlp_family_site_cs(first, last, C)) + return glu_site_specs(cfg, mlp_family_site_cs(first, last, C)) _QVDOWN_SITE_CS = ( @@ -163,10 +161,10 @@ def test_site_name_helpers(): canonical_site_cs((SiteC("layers.3.mlp.up_proj", 4), SiteC("layers.3.mlp.up_proj", 8))) -def test_llama_site_specs_dims(): +def test_glu_site_specs_dims(): cfg = _tiny_cfg() qd, kvd = cfg.n_head * cfg.head_dim, cfg.n_kv_head * cfg.head_dim - specs = llama_site_specs( + specs = glu_site_specs( cfg, canonical_site_cs( tuple(SiteC(site_name(2, kind), 4) for kind in ("q", "k", "v", "o", "gate", "down")) @@ -183,7 +181,7 @@ def dims(s: SiteSpec) -> tuple[int, int, int]: assert dims(by_name["layers.2.mlp.gate_proj"]) == (cfg.n_embd, cfg.n_intermediate, 4) assert dims(by_name["layers.2.mlp.down_proj"]) == (cfg.n_intermediate, cfg.n_embd, 4) with pytest.raises(AssertionError, match="canonical"): - llama_site_specs(cfg, tuple(reversed(mlp_family_site_cs(2, 2, 4)))) + glu_site_specs(cfg, tuple(reversed(mlp_family_site_cs(2, 2, 4)))) def test_masked_component_activations_pre_mask_and_matches_outputs(): @@ -261,7 +259,7 @@ def test_clean_path_and_masked_identity(first: int, last: int): def test_attention_sites_clean_and_masked_identity(): cfg = _tiny_cfg() - sites = llama_site_specs(cfg, _QVDOWN_SITE_CS) + sites = glu_site_specs(cfg, _QVDOWN_SITE_CS) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) b, t = 2, 16 @@ -330,7 +328,7 @@ def test_attention_sites_clean_and_masked_identity(): def test_o_site_masks_attention_output(): cfg = _tiny_cfg() o_site = "layers.4.self_attn.o_proj" - sites = llama_site_specs(cfg, (SiteC(o_site, 8),)) + sites = glu_site_specs(cfg, (SiteC(o_site, 8),)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) b, t = 2, 16 @@ -362,7 +360,7 @@ def test_step_trains_and_has_vpd_signature(site_cs: tuple[SiteC, ...]): cfg = _tiny_cfg() seq = 16 n_warmup = 2 - sites = llama_site_specs(cfg, site_cs) + sites = glu_site_specs(cfg, site_cs) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) ci_fn = _build_chunkwise_ci_fn(lm, jax.random.PRNGKey(2), n_blocks=2) @@ -468,7 +466,7 @@ def test_faith_warmup_decreases_faith(): def test_decomp_vu_shapes_fp32(): cfg = _tiny_cfg() - sites = llama_site_specs(cfg, _QVDOWN_SITE_CS) + sites = glu_site_specs(cfg, _QVDOWN_SITE_CS) vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) d, di = cfg.n_embd, cfg.n_intermediate qd, kvd = cfg.n_head * cfg.head_dim, cfg.n_kv_head * cfg.head_dim @@ -494,7 +492,7 @@ def test_fresh_pgd_adversary_step(): SiteC("layers.4.mlp.down_proj", 12), ) seq = 16 - sites = llama_site_specs(cfg, site_cs) + sites = glu_site_specs(cfg, site_cs) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0)) opt_ci = optax.adamw(1e-3, weight_decay=0.0) diff --git a/param_decomp/tests/test_llama_simple_mlp.py b/param_decomp/tests/test_llama_simple_mlp.py index 567e524ea..b23f7e744 100644 --- a/param_decomp/tests/test_llama_simple_mlp.py +++ b/param_decomp/tests/test_llama_simple_mlp.py @@ -37,7 +37,7 @@ from param_decomp.lm import DecomposedModel from param_decomp.recon import build_loss_terms from param_decomp.schedule import ScheduleConfig -from param_decomp.targets.llama8b import FrozenAttn +from param_decomp.targets.glu_transformer import FrozenAttn from param_decomp.targets.llama_simple_mlp import ( LlamaSimpleMLPConfig, SimpleMLPDecomposedModel, diff --git a/param_decomp/tests/test_qwen3_8b.py b/param_decomp/tests/test_qwen3_8b.py new file mode 100644 index 000000000..25b99d346 --- /dev/null +++ b/param_decomp/tests/test_qwen3_8b.py @@ -0,0 +1,233 @@ +"""CPU tests for the Qwen3 family target at a tiny config. + +Mirrors `test_llama8b.py` over the shared GLU-transformer machinery: the +`DecomposedModel` contract (clean == all-frozen masked forward, mask=1 identity, ablation +changes logits) and one full SPEC step — plus the family-specific pin that the QK-norm is +actually load-bearing in the forward. Direct HF parity lives in `tests/qwen3_hf_parity/`. +""" + +import equinox as eqx +import jax +import jax.numpy as jnp + +from param_decomp.components import SiteC, SiteSpec, init_decomp_vu +from param_decomp.targets.glu_transformer import ( + GLUConfig, + GLUDecomposedModel, + GLULayer, + build_decomposed_lm, + default_inv_freq, + glu_site_specs, + parse_site_name, +) +from param_decomp.targets.qwen3_8b import Qwen3FrozenAttn + + +def _tiny_cfg() -> GLUConfig: + """Qwen3-shaped tiny config: QK-norm attention, plain RoPE.""" + return GLUConfig( + vocab_size=64, + n_layer=8, + n_head=4, + n_kv_head=2, + n_embd=32, + n_intermediate=64, + rope_theta=1000000.0, + rms_norm_eps=1e-6, + max_position_embeddings=512, + ) + + +def _tiny_decomposed_qwen( + cfg: GLUConfig, sites: tuple[SiteSpec, ...], key: jax.Array +) -> GLUDecomposedModel: + """A tiny random Qwen3-family model — the CPU-test analog of + `load_decomposed_qwen3_from_hf` (`test_llama8b._tiny_decomposed_lm`'s sibling).""" + ks = iter(jax.random.split(key, 1024)) + d, di = cfg.n_embd, cfg.n_intermediate + qd, kvd = cfg.n_head * cfg.head_dim, cfg.n_kv_head * cfg.head_dim + + def n(shape: tuple[int, ...], s: float | None = None) -> jax.Array: + return jax.random.normal(next(ks), shape) * (s or d**-0.5) + + def fattn(): + return Qwen3FrozenAttn( + n((qd, d)), n((kvd, d)), n((kvd, d)), n((d, qd)), + cfg.n_head, cfg.n_kv_head, cfg.head_dim, cfg.n_rep, + # non-trivial norm weights (≈1) so a wrong/missing norm application shows + q_norm=1.0 + 0.1 * jax.random.normal(next(ks), (cfg.head_dim,)), + k_norm=1.0 + 0.1 * jax.random.normal(next(ks), (cfg.head_dim,)), + eps=cfg.rms_norm_eps, + ) # fmt: skip + + def layer(): + return GLULayer(jnp.ones((d,)), jnp.ones((d,)), fattn(), n((di, d)), n((di, d)), n((d, di))) + + return build_decomposed_lm( + embed=n((cfg.vocab_size, d), 0.02), + layers=[layer() for _ in range(cfg.n_layer)], + norm=jnp.ones((d,)), + lm_head=n((cfg.vocab_size, d), 0.02), + inv_freq=default_inv_freq(cfg.head_dim, cfg.rope_theta), + cfg=cfg, + sites=sites, + ) + + +_QVDOWN_SITE_CS = ( + SiteC("layers.4.self_attn.q_proj", 8), + SiteC("layers.4.self_attn.v_proj", 12), + SiteC("layers.4.mlp.down_proj", 8), +) +"""Attention + MLP sites on one layer with heterogeneous per-site C.""" + + +def test_clean_path_and_masked_identity(): + cfg = _tiny_cfg() + sites = glu_site_specs(cfg, _QVDOWN_SITE_CS) + lm = _tiny_decomposed_qwen(cfg, sites, jax.random.PRNGKey(0)) + vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) + b, t = 2, 16 + tokens = jax.random.randint(jax.random.PRNGKey(2), (b, t), 0, cfg.vocab_size) + + clean = lm.clean_output(tokens) + assert clean.shape == (b, t, cfg.vocab_size) + + # SPEC S2: a masked forward with NO live sites is the frozen path — bit-identical. + none_masked = lm.masked_output( + lm.prepare_compute_weights(vu), tokens, {}, {}, None, (), True, remat=False + ) + assert jnp.array_equal(clean, none_masked), "live=() must be the exact frozen path" + + # mask=1 identity through the QK-norm (V@U + (W − V@U); exact only in exact math). + names = lm.site_names + ones_masks = {s.name: jnp.ones((b, t, s.C)) for s in lm.sites} + ones_delta = {s: jnp.ones((b, t)) for s in names} + full = lm.masked_output( + lm.prepare_compute_weights(vu), tokens, ones_masks, ones_delta, None, names, True, + remat=False, + ) # fmt: skip + assert jnp.allclose(clean, full, atol=1e-4), "mask=1 identity drifted" + + # zero-mask + zero-delta on layer 4's decomposed sites must CHANGE the logits (q is + # live on the attention path ahead of QK-norm/RoPE/SDPA). + zero_mask = {s.name: jnp.zeros((b, t, s.C)) for s in lm.sites} + zero_delta = {s: jnp.zeros((b, t)) for s in names} + ablated = lm.masked_output( + lm.prepare_compute_weights(vu), tokens, zero_mask, zero_delta, None, names, True, + remat=False, + ) # fmt: skip + assert not jnp.allclose(clean, ablated, atol=1e-4), "ablating layer 4 did nothing" + + +def test_qk_norm_is_load_bearing(): + """The QK-norm actually enters the forward: scaling ONE layer's `q_norm` changes the + logits; the pre-projection q/k/v site inputs stay untouched (q/k sites decompose + BEFORE the norm — the masked site output feeds q_norm → RoPE → SDPA); o's site input + (the attention output) responds.""" + cfg = _tiny_cfg() + sites = glu_site_specs(cfg, _QVDOWN_SITE_CS) + lm = _tiny_decomposed_qwen(cfg, sites, jax.random.PRNGKey(0)) + tokens = jax.random.randint(jax.random.PRNGKey(2), (2, 16), 0, cfg.vocab_size) + + attn = lm.stacked.attn + assert isinstance(attn, Qwen3FrozenAttn) + assert attn.q_norm.shape == (cfg.n_layer, cfg.head_dim) + # scale ONLY layer 4's q_norm so the residual ENTERING layer 4 stays untouched + scaled = eqx.tree_at(lambda m: m.stacked.attn.q_norm, lm, attn.q_norm.at[4].mul(2.0)) + assert not jnp.allclose(lm.clean_output(tokens), scaled.clean_output(tokens), atol=1e-4) + + q_site = "layers.4.self_attn.q_proj" + taps = lm.read_activations(tokens, lm.site_names) + scaled_taps = scaled.read_activations(tokens, lm.site_names) + assert jnp.array_equal(taps[q_site], scaled_taps[q_site]) + o_site = "layers.4.self_attn.o_proj" + o_tap = lm.read_activations(tokens, (o_site,))[o_site] + o_tap_scaled = scaled.read_activations(tokens, (o_site,))[o_site] + assert not jnp.allclose(o_tap, o_tap_scaled) + + +def test_attn_pattern_applies_that_layers_qk_norm(): + """The target-owned attn-pattern recipe uses the site's LAYER norms: the same q/k + flats produce different patterns for two layers whose q_norm weights differ.""" + cfg = _tiny_cfg() + site_cs = ( + SiteC("layers.4.self_attn.q_proj", 8), + SiteC("layers.4.self_attn.k_proj", 8), + SiteC("layers.5.self_attn.q_proj", 8), + SiteC("layers.5.self_attn.k_proj", 8), + ) + sites = glu_site_specs(cfg, site_cs) + lm = _tiny_decomposed_qwen(cfg, sites, jax.random.PRNGKey(0)) + attn = lm.stacked.attn + assert isinstance(attn, Qwen3FrozenAttn) + lm = eqx.tree_at(lambda m: m.stacked.attn.q_norm, lm, attn.q_norm.at[5].mul(3.0)) + + b, t = 2, 9 + qd, kvd = cfg.n_head * cfg.head_dim, cfg.n_kv_head * cfg.head_dim + q = jax.random.normal(jax.random.PRNGKey(1), (b, t, qd)) + k = jax.random.normal(jax.random.PRNGKey(2), (b, t, kvd)) + p4 = lm.attn_pattern("layers.4.self_attn.q_proj", q, k) + p5 = lm.attn_pattern("layers.5.self_attn.q_proj", q, k) + assert p4.shape == (b, cfg.n_head, t, t) + assert not jnp.allclose(p4, p5) + assert parse_site_name("layers.5.self_attn.q_proj") == (5, "q") + + +def test_step_trains(): + """One full generic train step over the qwen tiny target — pins that the family plugs + into the engine (the step machinery itself is pinned in `test_llama8b.py`).""" + import optax + + from param_decomp.configs import ( + ChunkwiseSubsetReconLossConfig, + FaithfulnessLossConfig, + ImportanceMinimalityLossConfig, + UniformKSubsetRoutingConfig, + ) + from param_decomp.recon import build_loss_terms + from param_decomp.schedule import ScheduleConfig + from param_decomp.tests.test_llama8b import _build_chunkwise_ci_fn + from param_decomp.train import TrainState, make_train_step + + cfg = _tiny_cfg() + sites = glu_site_specs(cfg, _QVDOWN_SITE_CS) + lm = _tiny_decomposed_qwen(cfg, sites, jax.random.PRNGKey(0)) + vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) + ci_fn = _build_chunkwise_ci_fn(lm, jax.random.PRNGKey(2), n_blocks=1) + opt_vu = optax.adamw(1e-3, weight_decay=0.0) + opt_ci = optax.adamw(1e-3, weight_decay=0.0) + state = TrainState( + components=vu, ci_fn=ci_fn, + components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)), + ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)), + adversaries={}, + step=jnp.zeros((), jnp.int32), + ) # fmt: skip + loss_terms = build_loss_terms( + ( + FaithfulnessLossConfig(coeff=1e5), + ImportanceMinimalityLossConfig( + coeff=5e-6, + pnorm=ScheduleConfig(start_val=2.0, fn_type="linear", final_val_frac=0.2), + ), + ChunkwiseSubsetReconLossConfig( + routing=UniformKSubsetRoutingConfig(), coeff=0.5, sites_per_chunk=3, n_samples=1 + ), + ), + lm.site_names, + ) + step = make_train_step( + lm=lm, + losses=loss_terms, + components_optimizer=opt_vu, + ci_fn_optimizer=opt_ci, + total_steps=100, + remat_recon_forwards=True, + remat_ci_fn=False, + mesh=None, + ) + tokens = jax.random.randint(jax.random.PRNGKey(4), (2, 16), 0, cfg.vocab_size) + state, metrics = step(lm, state, tokens, jax.random.PRNGKey(100)) + assert all(jnp.isfinite(jnp.asarray(v)).all() for v in metrics.values()) + assert int(state.step) == 1 diff --git a/param_decomp/tests/test_sharding.py b/param_decomp/tests/test_sharding.py index cf4346435..40a6348e6 100644 --- a/param_decomp/tests/test_sharding.py +++ b/param_decomp/tests/test_sharding.py @@ -62,8 +62,8 @@ def test_jitted_sharded_inits_match_eager_values(): ) from param_decomp.components import SiteC, init_decomp_vu from param_decomp.configs import BSCScope, SCScope - from param_decomp.targets.llama8b import canonical_site_cs, llama_site_specs - from param_decomp.targets.llama8b_sharding import ( + from param_decomp.targets.glu_transformer import canonical_site_cs, glu_site_specs + from param_decomp.targets.glu_transformer_sharding import ( init_ci_fn_placed, init_decomp_vu_placed, init_sources_sharded, @@ -80,7 +80,7 @@ def test_jitted_sharded_inits_match_eager_values(): # layers.{2,3}.mlp.gate_proj share (d_in, d_out, C), so the stacked init has a REAL # multi-site V/U shape group (stack + unstack across sites, not just groups of one); # they also make a 3-site C group for the stacked sources init below. - sites = llama_site_specs( + sites = glu_site_specs( cfg, canonical_site_cs( ( diff --git a/param_decomp/tests/test_slow_eval.py b/param_decomp/tests/test_slow_eval.py index b0a44ecd4..6f23102ba 100644 --- a/param_decomp/tests/test_slow_eval.py +++ b/param_decomp/tests/test_slow_eval.py @@ -47,8 +47,8 @@ render_slow_eval_figures, resolve_permutation_metrics, ) -from param_decomp.targets.llama8b import ( - llama_site_specs, +from param_decomp.targets.glu_transformer import ( + glu_site_specs, mlp_family_site_cs, ) from param_decomp.tests.test_llama8b import ( @@ -77,7 +77,7 @@ def _build_ci_fn(lm: DecomposedModel, n_embd: int, key: jax.Array) -> CIFn: def _tiny_setup(threshold: float, density_heatmap_n_bins: int | None = None): cfg = _tiny_cfg() C = 8 - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, C)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) ci_fn = _build_ci_fn(lm, cfg.n_embd, jax.random.PRNGKey(2)) step = make_slow_eval_step(lm, threshold, density_heatmap_n_bins) @@ -308,7 +308,7 @@ def test_resolve_permutation_metrics_empty_when_unconfigured(): def _tiny_position_ci(): cfg = _tiny_cfg() - sites = llama_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) + sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, 8)) lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0)) ci_fn = _build_ci_fn(lm, cfg.n_embd, jax.random.PRNGKey(2)) residual = jax.random.randint(jax.random.PRNGKey(4), (3, 12), 0, cfg.vocab_size) diff --git a/param_decomp/tests/test_source_grad_mean.py b/param_decomp/tests/test_source_grad_mean.py index 71a20efdc..94d690ac3 100644 --- a/param_decomp/tests/test_source_grad_mean.py +++ b/param_decomp/tests/test_source_grad_mean.py @@ -42,8 +42,8 @@ from param_decomp.components import init_decomp_vu from param_decomp.losses import kl_per_position from param_decomp.sharding import hsdp_mesh, shard_batch -from param_decomp.targets.llama8b import ( - llama_site_specs, +from param_decomp.targets.glu_transformer import ( + glu_site_specs, mlp_family_site_cs, ) from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm @@ -56,7 +56,7 @@ def _source_grad(sharded: bool) -> dict[str, jax.Array]: reduction we are pinning. Returns one fp32 grad array per site.""" cfg = _tiny_cfg() C, seq, gbatch = 8, 16, 8 - sites = llama_site_specs(cfg, mlp_family_site_cs(3, 6, C)) + sites = glu_site_specs(cfg, mlp_family_site_cs(3, 6, C)) lm = _tiny_decomposed_lm(cfg, sites, random.PRNGKey(0)) vu = init_decomp_vu(sites, random.PRNGKey(1)) first_block = min(int(name.split(".")[1]) for name in lm.site_names) diff --git a/param_decomp/tools/theoretical_min_memory.py b/param_decomp/tools/theoretical_min_memory.py index 92021ff92..6e7a236c9 100644 --- a/param_decomp/tools/theoretical_min_memory.py +++ b/param_decomp/tools/theoretical_min_memory.py @@ -93,7 +93,7 @@ def main() -> None: trainable = vu + ci target = N_LAYERS * sum(d_in * d_out for d_in, d_out in KIND_DIMS.values()) # embed + lm_head are NOT decomposed and are REPLICATED (not ÷fsdp) on the frozen target - # (targets/llama8b.py:449-450) — full-resident per GPU. Llama-8B does not tie them. + # (targets/glu_transformer.py) — full-resident per GPU. Llama-8B does not tie them. embed_lmhead = 2 * VOCAB * D_MODEL GB = 1024**3 @@ -101,7 +101,7 @@ def main() -> None: opt = trainable * 12 / dp / GB # EXACT (sharding.py: master+m+v ÷N) vu_bf16 = vu * 2 / fsdp / GB # ÷fsdp resident compute weight (EXACT layout) ci_bf16 = ci * 2 / fsdp / GB # ÷fsdp resident compute weight (EXACT layout) - tgt_bf16 = target * 2 / fsdp / GB # EXACT: layer weights ÷fsdp (targets/llama8b.py:236-244) + tgt_bf16 = target * 2 / fsdp / GB # EXACT: layer weights ÷fsdp (targets/glu_transformer.py) embed_bf16 = embed_lmhead * 2 / GB # EXACT: embed+lm_head REPLICATED (full-resident) # ONE forward's activations (per-layer remat => residual carry stack + logits), per GPU. resid_stack = N_LAYERS * seq_per_gpu * SEQ * D_MODEL * 2 / GB # bf16 [L,b,t,d] carry diff --git a/param_decomp_lab/experiments/CLAUDE.md b/param_decomp_lab/experiments/CLAUDE.md index 246e41758..a0b9ce27f 100644 --- a/param_decomp_lab/experiments/CLAUDE.md +++ b/param_decomp_lab/experiments/CLAUDE.md @@ -88,8 +88,20 @@ target: kind: hf_weights_in_vendored # HF weights loaded into a vendored, componentizable arch model_class: param_decomp_lab.experiments.lm.vendored.llama_3_1.model.VendoredLlama model_name: meta-llama/Llama-3.1-8B + +# or +target: + spec: + kind: hf # Qwen3-8B-Base (same vendored JAX target + QK-norm) + model_class: transformers.Qwen3ForCausalLM + model_name: Qwen/Qwen3-8B-Base ``` +`kind: hf`/`hf_weights_in_vendored` model names must be in `experiments/lm/config.py::HF_MODEL_FAMILIES` +(Llama-3.1-8B, Qwen3-8B-Base) — anything else refuses at convert time. A Qwen3 run needs +a Qwen3-tokenized prestaged dataset (`prestage_tokenized` with +`--tokenizer_name Qwen/Qwen3-8B-Base`). + The JAX prediction tensor is always the final logits (there is no `output_extract` — it was a torch-era field, stripped on load for back-compat). The `model_class` strings are NOT imported by the JAX trainer — `param_decomp.built_run` only asserts the class-name diff --git a/param_decomp_lab/experiments/lm/config.py b/param_decomp_lab/experiments/lm/config.py index f49d3062f..20352aea6 100644 --- a/param_decomp_lab/experiments/lm/config.py +++ b/param_decomp_lab/experiments/lm/config.py @@ -8,6 +8,7 @@ `build_from_schema`; consumers that read a finished run dir call `load_run_dir_config`. """ +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import Annotated, Any, Literal @@ -40,8 +41,8 @@ StochasticAttnPatternsReconLossConfig, ) from param_decomp.recon import build_loss_terms -from param_decomp.targets import llama8b, llama_simple_mlp -from param_decomp.targets.llama8b import SITE_NAME_PATTERN, canonical_site_cs +from param_decomp.targets import glu_transformer, llama8b, llama_simple_mlp, qwen3_8b +from param_decomp.targets.glu_transformer import SITE_NAME_PATTERN, canonical_site_cs from param_decomp_lab.experiments.config import ( ExperimentConfig, assert_canonical_algorithm_config, @@ -145,17 +146,62 @@ class LMExperimentConfig(ExperimentConfig[LMTargetConfig, LMDataConfig]): pass +@dataclass(frozen=True) +class HFModelFamily: + """One vendored HF model family the LM composition can target: its arch config, its + HF loader (the family file's `load_decomposed_*_from_hf`), and the path-schema model + type consumers key on. The families live in `param_decomp/targets/{llama8b,qwen3_8b}.py` + over the shared `glu_transformer` machinery; this registry is the ONLY place a model + name selects a family.""" + + arch_config: Callable[[], glu_transformer.GLUArch] + load: Callable[..., glu_transformer.GLUDecomposedModel] + """`(model_name, cfg, sites, scan_unroll=..., gather_fp8=...)` — cfg is the family's + own arch-config type, so the common signature is erased here.""" + model_type: str + model_class: str + """The `target.spec.model_class` this family answers to (a stable identifier, never + imported — see experiments/CLAUDE.md).""" + + +HF_MODEL_FAMILIES: dict[str, HFModelFamily] = { + "meta-llama/Llama-3.1-8B": HFModelFamily( + llama8b.llama31_8b_config, + llama8b.load_decomposed_llama_from_hf, + "Llama", + "transformers.LlamaForCausalLM", + ), + "Qwen/Qwen3-8B-Base": HFModelFamily( + qwen3_8b.qwen3_8b_config, + qwen3_8b.load_decomposed_qwen3_from_hf, + "Qwen3", + "transformers.Qwen3ForCausalLM", + ), +} +"""The HF model names the LM composition implements. Anything else refuses loudly at +convert time — a new model gets an explicit family entry (config checked against its HF +config.json), never a silent guess.""" + + +def hf_model_family(model_name: str) -> HFModelFamily: + assert model_name in HF_MODEL_FAMILIES, ( + f"no vendored model family for {model_name!r}; supported: {sorted(HF_MODEL_FAMILIES)}" + ) + return HF_MODEL_FAMILIES[model_name] + + @dataclass(frozen=True) class TargetConfig: - """The Llama-3.1-8B HF target (`param_decomp.llama8b`).""" + """An HF GLU-transformer target (`model_name` must be in `HF_MODEL_FAMILIES` — + Llama-3.1-8B or Qwen3-8B-Base).""" model_name: str sites: tuple[SiteC, ...] """Decomposed sites with per-site C, in canonical order (`canonical_site_cs`).""" supported_weights_dtypes: frozenset[WeightsDtype] = frozenset({"bfloat16"}) - """Frozen-target weight dtypes the loader supports (`llama8b.py` is bf16-only: - `DT = jnp.bfloat16`). A config requesting a dtype outside this set is refused at + """Frozen-target weight dtypes the loader supports (the HF family loaders pass + bf16). A config requesting a dtype outside this set is refused at convert time — no silent downgrade (issue #727).""" @@ -230,9 +276,14 @@ def _resolve_target(cfg: LMExperimentConfig) -> AnyLMTargetConfig: match spec: case HFWeightsInVendored(): assert spec.model_class.rsplit(".", 1)[-1] == "VendoredLlama", spec.model_class + assert "Llama-3.1-8B" in spec.model_name, spec.model_name case HFTarget(): - assert spec.model_class == "transformers.LlamaForCausalLM", spec.model_class - assert "Llama-3.1-8B" in spec.model_name, spec.model_name + known_classes = {f.model_class for f in HF_MODEL_FAMILIES.values()} + assert spec.model_class in known_classes, spec.model_class + assert spec.model_class == hf_model_family(spec.model_name).model_class, ( + f"{spec.model_class!r} is not {spec.model_name!r}'s family" + ) + hf_model_family(spec.model_name) # refuse unknown model names here return TargetConfig(model_name=spec.model_name, sites=_site_cs(cfg)) case PretrainedTarget(): assert spec.model_class.rsplit(".", 1)[-1] == "LlamaSimpleMLP", spec.model_class @@ -251,7 +302,7 @@ def _block_of_site(target: AnyLMTargetConfig, site_name: str) -> int: grammar (`layers.{i}...` for llama8b, `h.{i}...` for LlamaSimpleMLP).""" match target: case TargetConfig(): - return llama8b.parse_site_name(site_name)[0] + return glu_transformer.parse_site_name(site_name)[0] case LlamaSimpleMLPTargetConfig(): return llama_simple_mlp.parse_site_name(site_name)[0] @@ -261,7 +312,7 @@ def _resolve_d_resid(target: AnyLMTargetConfig) -> int: each chunk reads one residual tap of this width.""" match target: case TargetConfig(): - return llama8b.llama31_8b_config().n_embd + return hf_model_family(target.model_name).arch_config().n_embd case LlamaSimpleMLPTargetConfig(): cache_dir = llama_simple_mlp.pretrain_cache_dir(target.pretrain_run_path) return llama_simple_mlp.load_model_config(cache_dir).n_embd diff --git a/param_decomp_lab/experiments/lm/load_run.py b/param_decomp_lab/experiments/lm/load_run.py index c064d8591..7657830a7 100644 --- a/param_decomp_lab/experiments/lm/load_run.py +++ b/param_decomp_lab/experiments/lm/load_run.py @@ -40,16 +40,13 @@ from param_decomp.run_state import build_optimizers, init_train_state from param_decomp.sharding import hsdp_mesh, place_via_shardings from param_decomp.targets import llama_simple_mlp -from param_decomp.targets.llama8b import ( - llama31_8b_config, - llama_site_specs, - load_decomposed_lm_from_hf, -) -from param_decomp.targets.llama8b_sharding import place_target +from param_decomp.targets.glu_transformer import glu_site_specs +from param_decomp.targets.glu_transformer_sharding import place_target from param_decomp.train import COMPUTE_DT, TrainState, cast_floating from param_decomp_lab.experiments.lm.config import ( LlamaSimpleMLPTargetConfig, TargetConfig, + hf_model_family, load_run_dir_config, ) @@ -68,7 +65,7 @@ def build_target(cfg: BuiltRun, mesh: jax.sharding.Mesh) -> tuple[DecomposedMode """`(lm, vocab_size)` for the run's target config. The `lm` (an `eqx.Module`) IS the frozen target — it carries the full model weights (embedding included) as fields and embeds its token input internally. SimpleMLP reads its local pretrain cache (no network); - llama8b reads the HF snapshot (frozen bf16 weights + fp32-compute, matching `run.py::main`). + the HF families read the HF snapshot (frozen bf16 weights + fp32-compute, matching `run.py::main`). LM-only: harvest/slow-eval over the toy (TMS/ResidMLP) targets is not wired — those validate via their in-loop target-CI metric in the lab provider, not this path.""" @@ -83,19 +80,20 @@ def build_target(cfg: BuiltRun, mesh: jax.sharding.Mesh) -> tuple[DecomposedMode lm = place_via_shardings(loaded_lm, loaded_lm.shardings(mesh)) return lm, simple_cfg.vocab_size case TargetConfig(): - llama_cfg = llama31_8b_config() - sites = llama_site_specs(llama_cfg, cfg.target.sites) + family = hf_model_family(cfg.target.model_name) + arch_cfg = family.arch_config() + sites = glu_site_specs(arch_cfg, cfg.target.sites) lm = place_target( - load_decomposed_lm_from_hf( + family.load( cfg.target.model_name, - llama_cfg, + arch_cfg, sites, scan_unroll=cfg.runtime.scan_unroll, gather_fp8=cfg.runtime.gather_fp8, ), mesh, ) - return lm, llama_cfg.vocab_size + return lm, arch_cfg.vocab_size case _: raise AssertionError(f"build_target is LM-only; got target {type(cfg.target).__name__}") @@ -242,11 +240,12 @@ def run_metadata(run_dir: Path) -> RunMetadata: layer_activation_sizes=[(s.name, s.C) for s in cfg.target.sites], ) case TargetConfig(): - llama_cfg = llama31_8b_config() + family = hf_model_family(cfg.target.model_name) + arch_cfg = family.arch_config() return RunMetadata( - model_type="Llama", - n_blocks=llama_cfg.n_layer, - vocab_size=llama_cfg.vocab_size, + model_type=family.model_type, + n_blocks=arch_cfg.n_layer, + vocab_size=arch_cfg.vocab_size, layer_activation_sizes=[(s.name, s.C) for s in cfg.target.sites], ) case _: diff --git a/param_decomp_lab/experiments/lm/run.py b/param_decomp_lab/experiments/lm/run.py index fb6ea9170..36bdc508a 100644 --- a/param_decomp_lab/experiments/lm/run.py +++ b/param_decomp_lab/experiments/lm/run.py @@ -49,7 +49,6 @@ ) from param_decomp.attn_patterns_eval import ( accumulate_attn_patterns, - attn_pattern_for, attn_patterns_log_entries, make_ci_attn_patterns_step, make_stochastic_attn_patterns_step, @@ -93,7 +92,7 @@ resolve_permutation_metrics, stochastic_hidden_acts_n_mask_samples, ) -from param_decomp.targets.llama8b import hf_snapshot_dir +from param_decomp.targets.glu_transformer import hf_snapshot_dir from param_decomp.train import TrainState from param_decomp_lab.experiments.lm.arithmetic_probe import build_arithmetic_probe from param_decomp_lab.experiments.lm.config import ( @@ -230,7 +229,7 @@ def _make_arithmetic_eval( ) from transformers import AutoTokenizer # heavy import; only the arith probe needs it in-job - # resolve from the SAME snapshot the weights load from (llama8b.hf_snapshot_dir honors + # resolve from the SAME snapshot the weights load from (hf_snapshot_dir honors # HF_HUB_CACHE / the shared cluster cache) — the default HF cache may be empty on cluster tokenizer = AutoTokenizer.from_pretrained( str(hf_snapshot_dir(target.model_name)), local_files_only=True @@ -379,14 +378,11 @@ def _make_lm_eval_fn( ) attn_steps: dict[str, Any] = {} if eval.attn_patterns is not None: - pattern_fn = attn_pattern_for(lm) if eval.attn_patterns.ci_masked: - attn_steps["CIMaskedAttnPatternsReconLoss"] = make_ci_attn_patterns_step( - lm, pattern_fn, co - ) + attn_steps["CIMaskedAttnPatternsReconLoss"] = make_ci_attn_patterns_step(lm, co) if eval.attn_patterns.stochastic: attn_steps["StochasticAttnPatternsReconLoss"] = make_stochastic_attn_patterns_step( - lm, pattern_fn, eval.attn_patterns.stochastic_n_mask_samples, co + lm, eval.attn_patterns.stochastic_n_mask_samples, co ) slow_eval_step = make_slow_eval_step( diff --git a/param_decomp_lab/topology/path_schemas.py b/param_decomp_lab/topology/path_schemas.py index b7127f565..35a0109d8 100644 --- a/param_decomp_lab/topology/path_schemas.py +++ b/param_decomp_lab/topology/path_schemas.py @@ -205,11 +205,24 @@ class _GPT2PathSchema(_PathSchema): unembed_path = "lm_head" +class _HFGLUPathSchema(_PathSchema): + """The raw-HF GLU-transformer site grammar (`layers.{i}.self_attn.q_proj`, …) the + Llama/Qwen3 `DecomposedModel` targets emit (`param_decomp.targets.glu_transformer`).""" + + embedding_path = "embed_tokens" + blocks = "layers" + attn = _SeparateAttnPathSchema(base="self_attn", q="q_proj", k="k_proj", v="v_proj", o="o_proj") + mlp = _GLUPathSchema(base="mlp", gate="gate_proj", up="up_proj", down="down_proj") + unembed_path = "lm_head" + + _MODEL_TYPE_PATH_SCHEMAS: dict[str, type[_PathSchema]] = { "LlamaSimple": _LlamaSimplePathSchema, "LlamaSimpleMLP": _LlamaSimpleMLPPathSchema, "GPT2Simple": _GPT2SimplePathSchema, "GPT2": _GPT2PathSchema, + "Llama": _HFGLUPathSchema, + "Qwen3": _HFGLUPathSchema, } diff --git a/pyproject.toml b/pyproject.toml index 80d6a3210..2246f690c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -136,6 +136,7 @@ exclude = [ "param_decomp/tests/equivalence/torch_reference.py", # torch-env script "param_decomp/tools/convert_llama_simple_mlp_checkpoint.py", # torch-env script "param_decomp/tests/simple_mlp_equivalence/gen_torch_fixtures.py", # torch-env script + "param_decomp/tests/qwen3_hf_parity/gen_hf_fixtures.py", # torch-env script "param_decomp/tests/stacked_parity/gen_stacked_fixtures.py", # base-branch-only script (pre-restructure API) ] stubPath = "typings" # Having type stubs for transformers shaves 10 seconds off basedpyright calls