diff --git a/param_decomp/components.py b/param_decomp/components.py index 2fac68c75..36397b2df 100644 --- a/param_decomp/components.py +++ b/param_decomp/components.py @@ -12,7 +12,7 @@ from collections.abc import Iterator from dataclasses import dataclass from functools import cache -from typing import ClassVar, Generic +from typing import ClassVar, Generic, Literal import equinox as eqx import jax @@ -72,6 +72,9 @@ def dequantize_fp8(q: Array, scale: Array) -> Array: VUShape = tuple[int, int, int] # (d_in, d_out, C) +# How `init_component_stacks` seeds V/U; the config schema (`PDConfig.weight_init`) reuses it. +WeightInit = Literal["kaiming", "coupled"] + # site name -> (shape group, slot on the group's stack axis); static, canonical site order SiteSlots = tuple[tuple[str, VUShape, int], ...] @@ -150,7 +153,7 @@ def group_lengths(self) -> dict[VUShape, int]: return lengths -def init_stack_arrays( +def _kaiming_stack_arrays( sites: tuple[SiteSpec, ...], key: Array ) -> dict[VUShape, tuple[Array, Array]]: """Seeded stack init: `{(d_in, d_out, C): (V [g, d_in, C], U [g, C, d_out])}`, vmapped @@ -169,6 +172,41 @@ def init_stack_arrays( return stacked +def _coupled_stack_arrays( + sites: tuple[SiteSpec, ...], target_weights: dict[str, Array], key: Array +) -> dict[VUShape, tuple[Array, Array]]: + """Coupled stack init: unit-norm seed on the narrow side, wide side its raw W-image. + + `d_in <= d_out`: `v_c ~ unit norm`, `U_c <- (W v_c)^T`; else `u_c ~ unit norm`, + `V_c <- W^T u_c`. No C-dependent rescale — components sit at W's natural scale and + the component sum equals W restricted to the seed span (`W V V^T` when V is seeded), + the delta carrying the complement. One key per site, split in site order; vmapped + per shape group like `_kaiming_stack_arrays` so the compiled init doesn't scale with + site count.""" + keys = jax.random.split(key, len(sites)) + site_index = {spec.name: idx for idx, spec in enumerate(sites)} + stacked: dict[VUShape, tuple[Array, Array]] = {} + for (d_in, d_out, c), specs in vu_shape_groups(sites).items(): + ws = jnp.stack([target_weights[spec.name] for spec in specs]).astype(jnp.float32) + assert ws.shape == (len(specs), d_out, d_in), (ws.shape, (d_in, d_out, c)) + ks = keys[jnp.array([site_index[spec.name] for spec in specs])] + if d_in <= d_out: + + def seed_v(k: Array, w: Array, s: tuple[int, int] = (d_in, c)) -> tuple[Array, Array]: + v = jax.random.normal(k, s) + v = v / jnp.linalg.norm(v, axis=0, keepdims=True) + return v, (w @ v).T + else: + + def seed_v(k: Array, w: Array, s: tuple[int, int] = (c, d_out)) -> tuple[Array, Array]: + u = jax.random.normal(k, s) + u = u / jnp.linalg.norm(u, axis=1, keepdims=True) + return w.T @ u.T, u + + stacked[(d_in, d_out, c)] = jax.vmap(seed_v)(ks, ws) + return stacked + + def component_stacks_from_sites(vu: dict[str, tuple[Array, Array]]) -> ComponentStacks: """Build the stacked `ComponentStacks` from a per-site `{name: (V, U)}` dict (site order = dict order). The explicit-arrays constructor for toys and tests; the trainer inits @@ -187,11 +225,25 @@ def component_stacks_from_sites(vu: dict[str, tuple[Array, Array]]) -> Component return ComponentStacks(stacks=stacks, site_slots=site_slots_for(sites)) -def init_component_stacks(sites: tuple[SiteSpec, ...], key: Array) -> ComponentStacks: - """Small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5) per site, built directly in - the stacked persistence layout; the weight-delta channel carries the faithfulness - residual at init (before faithfulness warmup).""" - return ComponentStacks(stacks=init_stack_arrays(sites, key), site_slots=site_slots_for(sites)) +def init_component_stacks( + sites: tuple[SiteSpec, ...], + key: Array, + weight_init: WeightInit = "kaiming", + target_weights: dict[str, Array] | None = None, +) -> ComponentStacks: + """Seeded V/U masters in the stacked persistence layout, per `weight_init`: + 'kaiming' — small random fp32 V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5), ignores W + (the weight-delta channel carries the faithfulness residual at init); 'coupled' — + `_coupled_stack_arrays`, seeded from `target_weights` (required there, refused + otherwise).""" + match weight_init: + case "kaiming": + assert target_weights is None, "kaiming init draws blind — no target_weights" + stacks = _kaiming_stack_arrays(sites, key) + case "coupled": + assert target_weights is not None, "coupled init seeds from W — pass target_weights" + stacks = _coupled_stack_arrays(sites, target_weights, key) + return ComponentStacks(stacks=stacks, site_slots=site_slots_for(sites)) def site_out( diff --git a/param_decomp/configs.py b/param_decomp/configs.py index eea12ab23..b4c363511 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -31,6 +31,7 @@ ) from param_decomp.base_config import BaseConfig, Probability +from param_decomp.components import WeightInit from param_decomp.schedule import ScheduleConfig # --------------------------------------------------------------------------- @@ -1041,6 +1042,14 @@ def _strip_removed_jax_unsupported_fields(cls, data: object) -> object: default=0, description="Random seed for reproducibility, including LM dataset shuffling.", ) + weight_init: WeightInit = Field( + default="kaiming", + description="How component V/U are initialized. 'kaiming': iid normal " + "(V ~ N(0, d_in^-0.5), U ~ N(0, C^-0.5)), ignores W. 'coupled': unit-norm seed on " + "the narrow side, wide side its raw W-image (U_c from (W v_c)^T when d_in <= d_out, " + "else V_c from W^T u_c) — W-natural scale, component sum ~ W on a rank-C random " + "subspace (the delta carries the complement).", + ) loss_metrics: list[AnyLossMetricConfig] = Field( default_factory=list, description=( diff --git a/param_decomp/model.py b/param_decomp/model.py index 8cd95d759..8f4b7dfdf 100644 --- a/param_decomp/model.py +++ b/param_decomp/model.py @@ -217,6 +217,10 @@ def masked_site_outputs( the recon grid, which stays KL-on-final-logits.""" ... + def target_weight(self, name: str) -> Float[Array, "d_out d_in"]: + """The frozen target weight W for site `name`, in its stored dtype.""" + ... + def weight_deltas(self, vu: ComponentStacks) -> dict[str, Float[Array, "d_out d_in"]]: """fp32 `W − V@U` per site from the fp32 master `vu`.""" ... diff --git a/param_decomp/run.py b/param_decomp/run.py index 6116dd524..d2681cf76 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -516,7 +516,8 @@ def run_decomposition_training( rules = placement_rules if is_main: audit = component_stacks_audit( - eqx.filter_eval_shape(_partial(init_component_stacks, model.sites), init_key), rules + eqx.filter_eval_shape(_partial(init_component_stacks, model.sites), init_key), + rules, ) print( rules.describe( diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index a0c898743..2ec8615f0 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -20,6 +20,7 @@ from param_decomp.adversary import PersistentAdversary, init_sources_adam_state from param_decomp.ci_fn import ChunkwiseTransformerCIArch, CIFnArch +from param_decomp.components import WeightInit from param_decomp.configs import ( AdamPGDConfig, AdamWOptimizerConfig, @@ -163,14 +164,23 @@ def init_decomposition( init_key: PRNGKeyArray, mesh: Mesh, rules: PlacementRules, + weight_init: WeightInit = "kaiming", ) -> Decomposition: """The trained-product half of `init_train_state`, factored out so a consumer can `jax.eval_shape` it to recover the saved `decomposition` item's tree structure - without building (or knowing about) the optimizers/adversaries.""" + without building (or knowing about) the optimizers/adversaries (the `weight_init` + default is safe there — both inits produce the same tree).""" ci_key = random.fold_in(init_key, 1) + match weight_init: + case "kaiming": + target_weights = None + case "coupled": + target_weights = {name: model.target_weight(name) for name in model.site_names} # V/U placement derives from the rules table; the CI fn still declares its own # per-leaf shardings (the staged placement migration hasn't reached it). - components = init_component_stacks_placed(model.sites, init_key, rules) + components = init_component_stacks_placed( + model.sites, init_key, rules, weight_init, target_weights + ) ci_fn = init_ci_fn_placed(ci_fn_arch, model.sites, ci_key, mesh) assert ci_fn.has_position_axis == model.has_position_axis, ( f"CI fn has_position_axis={ci_fn.has_position_axis} but model declares " @@ -195,7 +205,7 @@ def init_train_state( assert isinstance(positions, Positioned) == model.has_position_axis, ( f"{positions} does not match the model's has_position_axis={model.has_position_axis}" ) - decomposition = init_decomposition(model, ci_fn_arch, init_key, mesh, rules) + decomposition = init_decomposition(model, ci_fn_arch, init_key, mesh, rules, pd.weight_init) components, ci_fn = decomposition.components, decomposition.ci_fn losses = build_loss_terms(pd.loss_metrics, model.site_names) persistent = persistent_configs(losses.recon) diff --git a/param_decomp/targets/glu_transformer.py b/param_decomp/targets/glu_transformer.py index d285406a7..23a3c7594 100644 --- a/param_decomp/targets/glu_transformer.py +++ b/param_decomp/targets/glu_transformer.py @@ -1136,12 +1136,17 @@ def masked_component_activations( assert set(collect_activations) == set(live), (sorted(collect_activations), sorted(live)) return collect_activations + def target_weight(self, name: str) -> Array: + """The frozen W for site `name` (stored dtype): the site's slice of the stacked + per-layer weights.""" + layer, kind = parse_site_name(name) + return _frozen_site_weight(jax.tree.map(lambda a, li=layer: a[li], self.stacked), kind) + def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]: """fp32 `W − V@U` per site from fp32 masters (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) + W = self.target_weight(spec.name) V, U = vu.site(spec.name) out[spec.name] = ( W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T diff --git a/param_decomp/targets/glu_transformer_sharding.py b/param_decomp/targets/glu_transformer_sharding.py index 08b334f5b..ec5e47d87 100644 --- a/param_decomp/targets/glu_transformer_sharding.py +++ b/param_decomp/targets/glu_transformer_sharding.py @@ -63,6 +63,7 @@ from param_decomp.components import ( ComponentStacks, SiteSpec, + WeightInit, init_component_stacks, ) from param_decomp.configs import SourceShape @@ -89,15 +90,23 @@ def place_target(tgt: GLUDecomposedModel, mesh: Mesh) -> GLUDecomposedModel: def init_component_stacks_placed( - sites: tuple[SiteSpec, ...], key: PRNGKeyArray, rules: PlacementRules + sites: tuple[SiteSpec, ...], + key: PRNGKeyArray, + rules: PlacementRules, + weight_init: WeightInit = "kaiming", + target_weights: dict[str, Array] | None = None, ) -> ComponentStacks: - """Seeded V/U init placed by `component_stacks_shardings(_, rules)` (the run's placement - policy), values bit-identical to the retired per-site init (pinned by `test_sharding`). - One jit, 2×n_shapes sharded outputs — the persistence layout IS the stacked layout, so - the old two-stage stack-then-unstack fan-out (and its transient extra copy) is gone.""" - abstract = eqx.filter_eval_shape(partial(init_component_stacks, sites), key) - placement = component_stacks_shardings(abstract, rules) - return jax.jit(partial(init_component_stacks, sites), out_shardings=placement)(key) + """`init_component_stacks` placed by `component_stacks_shardings(_, rules)` (the run's + placement policy); kaiming values bit-identical to the retired per-site init (pinned + by `test_sharding`). One jit, 2×n_shapes sharded outputs — the persistence layout IS + the stacked layout, so no host-side full tree ever exists; `key`/`target_weights` ride + as jit ARGUMENTS (not closure constants) so the frozen W arrays are not baked into the + compiled init.""" + init = partial(init_component_stacks, sites, weight_init=weight_init) + placement = component_stacks_shardings( + eqx.filter_eval_shape(init, key, target_weights=target_weights), rules + ) + return jax.jit(init, out_shardings=placement)(key, target_weights=target_weights) def init_ci_fn_placed( diff --git a/param_decomp/targets/llama_simple_mlp.py b/param_decomp/targets/llama_simple_mlp.py index 408237bdc..3b663a653 100644 --- a/param_decomp/targets/llama_simple_mlp.py +++ b/param_decomp/targets/llama_simple_mlp.py @@ -488,12 +488,16 @@ def masked_site_outputs( assert set(collect) == set(live), (sorted(collect), sorted(live)) return collect + def target_weight(self, name: str) -> Array: + """The frozen W for site `name`, in its stored dtype.""" + layer_idx, kind = parse_site_name(name) + return _frozen_site_weight(self.layers[layer_idx], kind) + def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]: """fp32 `W − V@U` per site from fp32 masters (faithfulness input).""" out: dict[str, Array] = {} for spec in self.sites: - layer_idx, kind = parse_site_name(spec.name) - W = _frozen_site_weight(self.layers[layer_idx], kind) + W = self.target_weight(spec.name) V, U = vu.site(spec.name) out[spec.name] = ( W.astype(jnp.float32) - (V.astype(jnp.float32) @ U.astype(jnp.float32)).T diff --git a/param_decomp/tests/test_attn_patterns_eval.py b/param_decomp/tests/test_attn_patterns_eval.py index 460adc2cf..976b4a7a2 100644 --- a/param_decomp/tests/test_attn_patterns_eval.py +++ b/param_decomp/tests/test_attn_patterns_eval.py @@ -275,6 +275,10 @@ def masked_site_outputs( del vu, resid, masks, delta_masks, routes, live, has_delta raise AssertionError("positionless stub fn must not be called") + def target_weight(self, name: str) -> jax.Array: + del name + raise AssertionError("positionless stub fn must not be called") + def weight_deltas(self, vu: Any) -> dict[str, jax.Array]: del vu raise AssertionError("positionless stub fn must not be called") diff --git a/param_decomp/tests/test_component_init.py b/param_decomp/tests/test_component_init.py new file mode 100644 index 000000000..6f1b52cc2 --- /dev/null +++ b/param_decomp/tests/test_component_init.py @@ -0,0 +1,60 @@ +"""Tests for the `coupled` mode of `init_component_stacks`.""" + +import jax +import jax.numpy as jnp + +from param_decomp.components import SiteSpec, init_component_stacks + + +def _random_weight(d_out: int, d_in: int) -> jax.Array: + return jax.random.normal(jax.random.PRNGKey(0), (d_out, d_in)) + + +def _col_space_residual(u: jax.Array, w: jax.Array) -> jax.Array: + q_out, _, _ = jnp.linalg.svd(w, full_matrices=False) + return u - (u @ q_out) @ q_out.T + + +def _row_space_residual(v: jax.Array, w: jax.Array) -> jax.Array: + _, _, vh = jnp.linalg.svd(w, full_matrices=False) + q_in = vh.T + return v - q_in @ (q_in.T @ v) + + +def _init_single(w: jax.Array, c: int, seed: int) -> tuple[jax.Array, jax.Array]: + d_out, d_in = w.shape + sites = (SiteSpec(name="m", d_in=d_in, d_out=d_out, C=c),) + vu = init_component_stacks(sites, jax.random.PRNGKey(seed), "coupled", {"m": w}) + return vu.site("m") + + +def test_coupled_seeds_are_unit_norm_and_derived_side_is_raw_w_image(): + w = _random_weight(10, 6) # d_in < d_out: V seeded, U derived + V, U = _init_single(w, c=8, seed=0) + assert jnp.allclose(jnp.linalg.norm(V, axis=0), jnp.ones(8), atol=1e-5) + assert jnp.allclose(U, (w @ V).T, atol=1e-6) + assert jnp.abs(_col_space_residual(U, w)).max() < 1e-5 + + w_t = w.T # d_in > d_out: U seeded, V derived + V_t, U_t = _init_single(w_t, c=8, seed=0) + assert jnp.allclose(jnp.linalg.norm(U_t, axis=1), jnp.ones(8), atol=1e-5) + assert jnp.allclose(V_t, w_t.T @ U_t.T, atol=1e-6) + assert jnp.abs(_row_space_residual(V_t, w_t)).max() < 1e-5 + + +def test_coupled_stacked_sites_get_independent_draws_coupled_to_their_own_w(): + # Two sites sharing one (d_in, d_out, C) shape group stack on one axis; each slice + # must be coupled to ITS OWN W with an independent seed draw. + key = jax.random.PRNGKey(3) + w_a = jax.random.normal(jax.random.fold_in(key, 0), (10, 6)) + w_b = jax.random.normal(jax.random.fold_in(key, 1), (10, 6)) + sites = ( + SiteSpec(name="a", d_in=6, d_out=10, C=8), + SiteSpec(name="b", d_in=6, d_out=10, C=8), + ) + vu = init_component_stacks(sites, jax.random.PRNGKey(0), "coupled", {"a": w_a, "b": w_b}) + V_a, U_a = vu.site("a") + V_b, U_b = vu.site("b") + assert not jnp.allclose(V_a, V_b) + assert jnp.allclose(U_a, (w_a @ V_a).T, atol=1e-6) + assert jnp.allclose(U_b, (w_b @ V_b).T, atol=1e-6) diff --git a/param_decomp/tests/test_eval.py b/param_decomp/tests/test_eval.py index 9fd5e790a..9327e18b5 100644 --- a/param_decomp/tests/test_eval.py +++ b/param_decomp/tests/test_eval.py @@ -131,6 +131,10 @@ def masked_site_outputs( del vu, resid, masks, delta_masks, routes, live, has_delta raise AssertionError("positionless stub fn must not be called") + def target_weight(self, name: str) -> jax.Array: + del name + raise AssertionError("positionless stub fn must not be called") + def weight_deltas(self, vu: Any) -> dict[str, jax.Array]: del vu raise AssertionError("positionless stub fn must not be called") diff --git a/param_decomp/tests/test_generic_model_io.py b/param_decomp/tests/test_generic_model_io.py index 281f25d46..f87fca0e7 100644 --- a/param_decomp/tests/test_generic_model_io.py +++ b/param_decomp/tests/test_generic_model_io.py @@ -166,6 +166,10 @@ def masked_site_outputs( hidden = hidden + delta_masks[SITE][..., None] * (resid @ delta.T) return {SITE: hidden} + def target_weight(self, name: str) -> Array: + assert name == SITE, name + return self.W + def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]: V, U = vu.site(SITE) return { diff --git a/param_decomp_lab/experiments/resid_mlp/model.py b/param_decomp_lab/experiments/resid_mlp/model.py index baffcb854..2854918d1 100644 --- a/param_decomp_lab/experiments/resid_mlp/model.py +++ b/param_decomp_lab/experiments/resid_mlp/model.py @@ -457,6 +457,9 @@ def masked_site_outputs( self.target, prepared, resid, masks, delta_masks, routes, live, has_delta ) + def target_weight(self, name: str) -> Array: + return _frozen_site_weight(self.target, name) + def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]: return weight_deltas_fp32(self.target, vu, self.sites) diff --git a/param_decomp_lab/experiments/tms/model.py b/param_decomp_lab/experiments/tms/model.py index 615b76822..fdbdbffad 100644 --- a/param_decomp_lab/experiments/tms/model.py +++ b/param_decomp_lab/experiments/tms/model.py @@ -389,6 +389,9 @@ def masked_site_outputs( self.target, prepared, resid, masks, delta_masks, routes, live, has_delta ) + def target_weight(self, name: str) -> Array: + return _frozen_site_weight(self.target, name) + def weight_deltas(self, vu: ComponentStacks) -> dict[str, Array]: return weight_deltas_fp32(self.target, vu, self.sites)