diff --git a/tests/integration/model_bridge/test_pretrain_adapter.py b/tests/integration/model_bridge/test_pretrain_adapter.py new file mode 100644 index 000000000..5878532e3 --- /dev/null +++ b/tests/integration/model_bridge/test_pretrain_adapter.py @@ -0,0 +1,312 @@ +"""Integration tests for PretrainArchitectureAdapter -- numerical parity +against a real (not structural-mock) reference model. + +Uses `tests/mocks/tiny_pretrain_model.py`'s `TinyPretrainModel`, which +implements genuine adjacent-pair RoPE and RMSNorm math (unlike +`_pretrain_mocks.py`'s unit-test mocks, which stub RoPE with constant +cos/sin and only need to get attribute names and call-signatures right). + +Covers: source/bridge logit parity (dense, full-MoE, mixed dense/MoE, +untied embeddings); that this file's fixtures supply a `cfg` truthfully +describing the wrapped model's head/vocab dimensions (a fixture property, +not something `build_pretrain_bridge` itself enforces -- see +`TestIntegrationFixturesSupplyATruthfulConfig`); residual-stream +decomposition via `run_with_cache`; a forward-hook intervention +cross-checked against an independent way of expressing the same edit; and +float64 construction-time preservation. + +Note on scope: since the adapter delegates the whole forward to the +source model rather than reimplementing attention/RoPE, exact logit +parity mainly demonstrates that wrapping and output normalization don't +alter an unhooked forward, not that the RoPE convention itself is +correct (the same implementation runs on both sides of the comparison). +`TestResidualStreamDecomposition` and `TestHookIntervention` are the +stronger evidence for correct intervention points. +""" + +from __future__ import annotations + +import torch + +from tests.mocks.tiny_pretrain_model import TinyPretrainModel +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.supported_architectures.pretrain import ( + ARCHITECTURE_NAME, + build_pretrain_bridge, +) + + +def _make_cfg( + *, + d_model: int, + n_layers: int, + n_heads: int = 2, + d_vocab: int = 64, +) -> TransformerBridgeConfig: + """Build a bridge config matching this file's TinyPretrainModel + fixtures (n_heads=2, vocab_size=64 by default). See + `TestIntegrationFixturesSupplyATruthfulConfig` for why this matching + matters and what happens when it doesn't.""" + return TransformerBridgeConfig( + d_model=d_model, + d_head=d_model // n_heads, + n_layers=n_layers, + n_ctx=128, + n_heads=n_heads, + d_vocab=d_vocab, + d_mlp=d_model * 2, + architecture=ARCHITECTURE_NAME, + ) + + +class TestSourceBridgeLogitParity: + """Bridge output must equal the source model's own `"logits"` output + exactly -- not within tolerance. The adapter never translates weights + into a second parameter layout, so there is no floating-point drift + to tolerate; any mismatch here is a real bug, not numerical noise.""" + + def test_dense_model_logit_parity_exact(self) -> None: + torch.manual_seed(0) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=2, d_ff=32, vocab_size=64) + model.eval() + tokens = torch.randint(0, 64, (1, 5)) + with torch.no_grad(): + source_logits = model(tokens)["logits"] + + bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2)) + with torch.no_grad(): + bridge_logits = bridge(tokens) + + torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0) + + def test_full_moe_model_logit_parity_exact(self) -> None: + """Every block's MLP is MoE (moe_layer_indices covers all layers).""" + torch.manual_seed(0) + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=2, + d_ff=32, + vocab_size=64, + n_experts=4, + top_k=2, + moe_layer_indices=frozenset({0, 1}), + ) + model.eval() + tokens = torch.randint(0, 64, (1, 5)) + with torch.no_grad(): + source_logits = model(tokens)["logits"] + + bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2)) + with torch.no_grad(): + bridge_logits = bridge(tokens) + + torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0) + + def test_mixed_dense_moe_model_logit_parity_exact(self) -> None: + """One block dense, one block MoE -- the dispatcher must route each + block's MLP correctly within a single model, not just across + separately-constructed all-dense/all-MoE models.""" + torch.manual_seed(0) + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=2, + d_ff=32, + vocab_size=64, + n_experts=4, + top_k=2, + moe_layer_indices=frozenset({1}), + ) + model.eval() + tokens = torch.randint(0, 64, (1, 5)) + with torch.no_grad(): + source_logits = model(tokens)["logits"] + + bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2)) + with torch.no_grad(): + bridge_logits = bridge(tokens) + + torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0) + + def test_untied_embeddings_logit_parity_exact(self) -> None: + torch.manual_seed(0) + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + tie_embeddings=False, + ) + model.eval() + tokens = torch.randint(0, 64, (1, 5)) + with torch.no_grad(): + source_logits = model(tokens)["logits"] + + bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=1)) + with torch.no_grad(): + bridge_logits = bridge(tokens) + + torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0) + + +class TestIntegrationFixturesSupplyATruthfulConfig: + """Integration fixtures must supply a config matching the source + model. `build_pretrain_bridge` does not itself validate this -- `cfg` + is caller-owned, per ordinary TransformerBridge convention -- so a + mismatch shows up only in `bridge.cfg`, silently, never in the + logits (see the second test).""" + + def test_cfg_matches_model_head_and_vocab_dimensions(self) -> None: + torch.manual_seed(0) + n_heads = 2 + vocab_size = 64 + model = TinyPretrainModel( + d_model=16, n_heads=n_heads, n_layers=2, d_ff=32, vocab_size=vocab_size + ) + model.eval() + cfg = _make_cfg(d_model=16, n_layers=2, n_heads=n_heads, d_vocab=vocab_size) + bridge = build_pretrain_bridge(model, cfg) + + assert bridge.cfg.n_heads == n_heads + assert bridge.cfg.d_head == 16 // n_heads + assert bridge.cfg.d_vocab == model.embed.num_embeddings + + def test_build_pretrain_bridge_does_not_validate_cfg_against_model(self) -> None: + """A cfg describing a different architecture than the wrapped + model still builds and still produces exact logit parity + (delegation never reads cfg) -- only bridge.cfg is wrong. + + Current framework convention; not an adapter guarantee. This test + documents present behavior, not a requirement -- if + build_pretrain_bridge later gains cfg-vs-model validation, this + test should be updated (or removed) rather than treated as a + regression to preserve.""" + torch.manual_seed(0) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + model.eval() + tokens = torch.randint(0, 64, (1, 5)) + with torch.no_grad(): + source_logits = model(tokens)["logits"] + + wrong_cfg = _make_cfg(d_model=16, n_layers=1, n_heads=4, d_vocab=256) + bridge = build_pretrain_bridge(model, wrong_cfg) + with torch.no_grad(): + bridge_logits = bridge(tokens) + + torch.testing.assert_close(bridge_logits, source_logits, atol=0, rtol=0) + assert bridge.cfg.n_heads != model.blocks[0].attn.n_heads + assert bridge.cfg.d_vocab != model.embed.num_embeddings + + +class TestResidualStreamDecomposition: + """`run_with_cache`'s resid hooks must decompose the way the residual + stream actually works, not just be present under the right names + (attribute/key presence is already covered by the structural unit + tests) -- resid_pre + attn_out == resid_mid, resid_mid + mlp_out == + resid_post, and one block's resid_post is the next block's resid_pre.""" + + def test_resid_stream_decomposes_exactly_via_run_with_cache(self) -> None: + torch.manual_seed(0) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=2, d_ff=32, vocab_size=64) + model.eval() + bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2)) + tokens = torch.randint(0, 64, (1, 5)) + + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens) + + torch.testing.assert_close( + cache["blocks.0.hook_resid_pre"], cache["hook_embed"], atol=0, rtol=0 + ) + torch.testing.assert_close( + cache["blocks.0.hook_resid_mid"], + cache["blocks.0.hook_resid_pre"] + cache["blocks.0.hook_attn_out"], + ) + torch.testing.assert_close( + cache["blocks.0.hook_resid_post"], + cache["blocks.0.hook_resid_mid"] + cache["blocks.0.hook_mlp_out"], + ) + torch.testing.assert_close( + cache["blocks.1.hook_resid_pre"], + cache["blocks.0.hook_resid_post"], + atol=0, + rtol=0, + ) + torch.testing.assert_close( + cache["blocks.1.hook_resid_post"], + cache["blocks.1.hook_resid_mid"] + cache["blocks.1.hook_mlp_out"], + ) + + +class TestHookIntervention: + """A forward-hook edit at `blocks.{i}.mlp.hook_out` must land at the + point the residual-stream decomposition says it should -- cross- + checked against an independently-expressed version of the same edit + (forcing resid_post to equal resid_mid), rather than only checking + that *some* change occurred.""" + + def test_mlp_hook_ablation_matches_independent_resid_post_override(self) -> None: + torch.manual_seed(0) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=2, d_ff=32, vocab_size=64) + model.eval() + bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=2)) + tokens = torch.randint(0, 64, (1, 5)) + + with torch.no_grad(): + baseline_logits, baseline_cache = bridge.run_with_cache(tokens) + resid_mid0 = baseline_cache["blocks.0.hook_resid_mid"].clone() + + def zero_mlp_out(value: torch.Tensor, hook) -> torch.Tensor: + return torch.zeros_like(value) + + def force_resid_post_to_resid_mid(value: torch.Tensor, hook) -> torch.Tensor: + return resid_mid0 + + with torch.no_grad(): + via_mlp_ablation = bridge.run_with_hooks( + tokens, fwd_hooks=[("blocks.0.mlp.hook_out", zero_mlp_out)] + ) + via_resid_override = bridge.run_with_hooks( + tokens, + fwd_hooks=[("blocks.0.hook_resid_post", force_resid_post_to_resid_mid)], + ) + + # Sanity: the intervention actually changed something. + assert not torch.equal(via_mlp_ablation, baseline_logits) + # The real check: two different-looking edits that are + # mathematically the same edit must produce identical output. + torch.testing.assert_close(via_mlp_ablation, via_resid_override, atol=0, rtol=0) + + +class TestDtypePreservation: + """float64 is checked for preservation through construction only, not + full numerical parity at that dtype -- RoPE's cos/sin computation + involves trig functions whose float32-vs-float64 rounding can + legitimately differ from the bridge's own dtype-casting path in ways + unrelated to correctness, so exact-equality parity isn't a meaningful + check at non-default dtypes the way it is at float32.""" + + def test_float64_construction_preserves_dtype_and_tied_weights(self) -> None: + torch.manual_seed(0) + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + tie_embeddings=True, + ).to(torch.float64) + assert model.lm_head.weight is model.embed.weight + + bridge = build_pretrain_bridge(model, _make_cfg(d_model=16, n_layers=1)) + + assert next(model.parameters()).dtype == torch.float64 + assert model.lm_head.weight is model.embed.weight + + tokens = torch.randint(0, 64, (1, 4)) + with torch.no_grad(): + output = bridge(tokens) + assert output.dtype == torch.float64 + assert torch.isfinite(output).all() diff --git a/tests/mocks/tiny_pretrain_model.py b/tests/mocks/tiny_pretrain_model.py new file mode 100644 index 000000000..2c3d9434c --- /dev/null +++ b/tests/mocks/tiny_pretrain_model.py @@ -0,0 +1,193 @@ +"""A tiny, self-contained decoder-only reference model (RoPE, RMSNorm, +gated SwiGLU MLP, optional MoE) used only by the pretrain-adapter +integration tests, to exercise real numerical parity rather than a purely +structural mock. + +Deliberately uses the *adjacent-pair* RoPE convention +(`[x0, x1] -> [-x1, x0]`), not HuggingFace's rotate-half (contiguous-half) +convention -- this is the concrete case `PretrainArchitectureAdapter`'s +docstring is about: an existing HF-oriented attention bridge would silently +apply the wrong rotation to a model using this convention. + +Also deliberately computes RoPE `cos`/`sin` once per forward pass and +passes them into every block call as extra positional args +(`block(x, cos, sin)`), matching the target architecture's convention. +This matters beyond efficiency: `BlockBridge` wraps a bare-tensor output +in a 1-tuple for single-positional-argument "standalone hidden_states" +calls (an HF-compatibility convention). A loop calling `block(x)` alone +would need to unwrap `[0]` from each result -- modifying the source +model's forward specifically because it's being bridged. Passing +`cos`/`sin` positionally avoids that, so the forward loop needs no +changes at all. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _rotate_adjacent_pairs(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., 0::2] + x2 = x[..., 1::2] + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + +def _apply_rotary_adjacent_pairs( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + cos = cos.repeat_interleave(2, dim=-1) + sin = sin.repeat_interleave(2, dim=-1) + return x * cos + _rotate_adjacent_pairs(x) * sin + + +class RMSNorm(nn.Module): + def __init__(self, d_model: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(d_model)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + return x * self.weight + + +class Attention(nn.Module): + """Receives precomputed `cos`/`sin` rather than computing them + internally -- matches the real target architecture, where rotary + tables are computed once per forward pass and shared across blocks.""" + + def __init__(self, d_model: int, n_heads: int): + super().__init__() + self.n_heads = n_heads + self.d_head = d_model // n_heads + self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) + self.o = nn.Linear(d_model, d_model, bias=False) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + b, s, d = x.shape + qkv = self.qkv(x).view(b, s, 3, self.n_heads, self.d_head).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # each: [b, n_heads, s, d_head] + + q = _apply_rotary_adjacent_pairs(q, cos, sin) + k = _apply_rotary_adjacent_pairs(k, cos, sin) + + attn = torch.softmax( + (q @ k.transpose(-2, -1)) / (self.d_head**0.5) + + torch.triu(torch.full((s, s), float("-inf"), device=x.device), diagonal=1), + dim=-1, + ) + out = (attn @ v).transpose(1, 2).reshape(b, s, d) + return self.o(out) + + +class DenseMLP(nn.Module): + def __init__(self, d_model: int, d_ff: int): + super().__init__() + self.gate = nn.Linear(d_model, d_ff, bias=False) + self.up = nn.Linear(d_model, d_ff, bias=False) + self.down = nn.Linear(d_ff, d_model, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down(F.silu(self.gate(x)) * self.up(x)) + + +class MoEMLP(nn.Module): + def __init__(self, d_model: int, d_ff: int, n_experts: int, top_k: int): + super().__init__() + self.top_k = top_k + self.router = nn.Linear(d_model, n_experts, bias=False) + self.experts = nn.ModuleList([DenseMLP(d_model, d_ff) for _ in range(n_experts)]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + weights = torch.softmax(self.router(x), dim=-1) + topk_weights, topk_idx = weights.topk(self.top_k, dim=-1) + out = torch.zeros_like(x) + for e in range(len(self.experts)): + mask = topk_idx == e + if not mask.any(): + continue + gate = (mask.float() * topk_weights).sum(dim=-1, keepdim=True) + out = out + gate * self.experts[e](x) + return out + + +class Block(nn.Module): + """Takes `cos`/`sin` as extra positional args alongside the hidden + state -- see the module docstring for why this matters beyond RoPE + plumbing.""" + + def __init__(self, d_model: int, n_heads: int, d_ff: int, mlp: nn.Module): + super().__init__() + self.norm1 = RMSNorm(d_model) + self.attn = Attention(d_model, n_heads) + self.norm2 = RMSNorm(d_model) + self.mlp = mlp + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x), cos, sin) + x = x + self.mlp(self.norm2(x)) + return x + + +class TinyPretrainModel(nn.Module): + """Full tiny reference model for numerical-parity integration tests. + Matches `PretrainArchitectureAdapter`'s expected attribute paths: + `embed`, `blocks`, `norm_f`, `lm_head`. + """ + + def __init__( + self, + d_model: int = 32, + n_heads: int = 4, + n_layers: int = 2, + d_ff: int = 64, + vocab_size: int = 256, + n_experts: int = 4, + top_k: int = 2, + moe_layer_indices: frozenset[int] = frozenset(), + tie_embeddings: bool = True, + rope_base: float = 10000.0, + ): + super().__init__() + self.embed = nn.Embedding(vocab_size, d_model) + self.blocks = nn.ModuleList( + [ + Block( + d_model, + n_heads, + d_ff, + mlp=( + MoEMLP(d_model, d_ff, n_experts, top_k) + if i in moe_layer_indices + else DenseMLP(d_model, d_ff) + ), + ) + for i in range(n_layers) + ] + ) + self.norm_f = RMSNorm(d_model) + self.lm_head = nn.Linear(d_model, vocab_size, bias=False) + if tie_embeddings: + self.lm_head.weight = self.embed.weight + + d_head = d_model // n_heads + inv_freq = 1.0 / (rope_base ** (torch.arange(0, d_head, 2).float() / d_head)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, tokens: torch.Tensor) -> dict: + b, s = tokens.shape + x = self.embed(tokens) + + # Computed once per forward pass, passed into every block -- the + # real target architecture's convention (see module docstring). + pos = torch.arange(s, device=tokens.device, dtype=self.inv_freq.dtype) + freqs = torch.outer(pos, self.inv_freq) + cos, sin = freqs.cos(), freqs.sin() + + for block in self.blocks: + x = block(x, cos, sin) + x = self.norm_f(x) + return {"logits": self.lm_head(x)} diff --git a/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py b/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py new file mode 100644 index 000000000..3b644497b --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_moe_bridge_tuple_output.py @@ -0,0 +1,88 @@ +"""Tests for MoEBridge tuple-output validation and router-score hooks. + +Unlike the pretrain-adapter's TestDenseOrMoEFeedForwardBridgeTupleOutput +(which patches _delegate.forward and therefore exercises +DenseOrMoEFeedForwardBridge.forward in isolation), these tests call +MoEBridge.forward() directly via a stub original_component, so they cover +MoEBridge's own tuple-validation contract rather than the dispatcher's. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from transformer_lens.model_bridge.generalized_components import MoEBridge + + +class _StubMoEModule(nn.Module): + """Minimal nn.Module whose forward is swapped in per-test, standing in + for a real MoE layer wrapped by MoEBridge.""" + + def __init__(self, fake_forward): + super().__init__() + # Give the module at least one parameter so + # `next(self.parameters())` in MoEBridge.forward doesn't raise + # StopIteration and skip the dtype-cast branch entirely. + self.dummy = nn.Linear(4, 4) + self._fake_forward = fake_forward + + def forward(self, *args, **kwargs): + return self._fake_forward(*args, **kwargs) + + +def _bridge_with_stub(fake_forward) -> MoEBridge: + bridge = MoEBridge(name="mlp") + bridge.set_original_component(_StubMoEModule(fake_forward)) + return bridge + + +class TestMoEBridgeTupleOutput: + def test_empty_tuple_raises_clear_type_error(self) -> None: + bridge = _bridge_with_stub(lambda *a, **kw: ()) + with pytest.raises(TypeError, match="torch.Tensor"): + bridge(torch.ones(1, 3, 4)) + + def test_non_tensor_metadata_preserved_without_router_hook(self) -> None: + bridge = _bridge_with_stub(lambda *a, **kw: (torch.zeros(1, 3, 4), "not_router_scores")) + fired = {"called": False} + + def mark_fired(value, hook): + fired["called"] = True + + bridge.hook_router_scores.add_hook(mark_fired) + + output = bridge(torch.ones(1, 3, 4)) + + assert isinstance(output, tuple) + assert output[1] == "not_router_scores" + torch.testing.assert_close(output[0], torch.zeros(1, 3, 4)) + assert fired["called"] is False + + def test_tensor_router_scores_still_fire_hook(self) -> None: + router_scores = torch.rand(1, 3, 8) + bridge = _bridge_with_stub(lambda *a, **kw: (torch.zeros(1, 3, 4), router_scores)) + captured = {} + + def capture_router_scores(value, hook): + captured["router_scores"] = value + return value + + bridge.hook_router_scores.add_hook(capture_router_scores) + + output = bridge(torch.ones(1, 3, 4)) + + assert isinstance(output, tuple) + torch.testing.assert_close(output[1], router_scores) + assert "router_scores" in captured + torch.testing.assert_close(captured["router_scores"], router_scores) + + def test_non_tensor_first_element_raises_clear_type_error(self) -> None: + """Exercises the output[0]-validation added to MoEBridge itself + (not just the dispatcher) -- a malformed first element should + raise a clear TypeError here rather than failing inside + HookPoint's own type checking.""" + bridge = _bridge_with_stub(lambda *a, **kw: ("not_a_tensor", torch.zeros(1, 3, 4))) + with pytest.raises(TypeError, match="torch.Tensor"): + bridge(torch.ones(1, 3, 4)) diff --git a/tests/unit/model_bridge/supported_architectures/_pretrain_mocks.py b/tests/unit/model_bridge/supported_architectures/_pretrain_mocks.py new file mode 100644 index 000000000..28c4f94f4 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/_pretrain_mocks.py @@ -0,0 +1,214 @@ +"""Shared mock modules and fixtures for PretrainArchitectureAdapter and +PretrainModelContainer unit tests. Not a test file itself (no `test_` +prefix, so pytest won't collect it) -- imported by both +test_pretrain_adapter.py and test_pretrain_model_container.py. + +Fully self-contained: builds tiny mock `nn.Module`s that reproduce the +relevant module tree (embed / blocks / norm1 / attn / norm2 / mlp / norm_f / +lm_head) rather than depending on any external training-framework package. +No weight loading, no network access. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig +from transformer_lens.model_bridge.supported_architectures.pretrain import ( + ARCHITECTURE_NAME, +) + +# -------------------------------------------------------------------------- +# Tiny mock module tree, matching the real target architecture's +# block-calling convention (cos/sin as extra positional args -- see +# tiny_pretrain_model.py's module docstring). Real RoPE/RMSNorm math isn't +# needed here (covered by integration parity tests); only correct +# attribute names/shapes/call-signatures for component_mapping to wire up. +# -------------------------------------------------------------------------- + + +class TinyRMSNorm(nn.Module): + def __init__(self, d_model: int): + super().__init__() + self.weight = nn.Parameter(torch.ones(d_model)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * self.weight + + +class TinyAttention(nn.Module): + """Opaque attention stand-in: linear-in, linear-out, no real RoPE math. + Takes (and ignores) cos/sin positionally, matching the real target + architecture's block-calling convention.""" + + def __init__(self, d_model: int, n_heads: int): + super().__init__() + self.n_heads = n_heads + self.qkv = nn.Linear(d_model, 3 * d_model) + self.out = nn.Linear(d_model, d_model) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + b, s, d = x.shape + qkv = self.qkv(x) + q, k, v = qkv.chunk(3, dim=-1) + attn = torch.softmax(q @ k.transpose(-2, -1) / (d**0.5), dim=-1) + return self.out(attn @ v) + + +class TinyDenseMLP(nn.Module): + def __init__(self, d_model: int, d_ff: int): + super().__init__() + self.gate = nn.Linear(d_model, d_ff, bias=False) + self.up = nn.Linear(d_model, d_ff, bias=False) + self.down = nn.Linear(d_ff, d_model, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down(torch.nn.functional.silu(self.gate(x)) * self.up(x)) + + +class TinyMoE(nn.Module): + def __init__(self, d_model: int, d_ff: int, n_experts: int, top_k: int): + super().__init__() + self.top_k = top_k + self.router = nn.Linear(d_model, n_experts, bias=False) + self.experts = nn.ModuleList([TinyDenseMLP(d_model, d_ff) for _ in range(n_experts)]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + weights = torch.softmax(self.router(x), dim=-1) + topk_weights, topk_idx = weights.topk(self.top_k, dim=-1) + out = torch.zeros_like(x) + for expert_idx in range(len(self.experts)): + mask = (topk_idx == expert_idx).any(dim=-1, keepdim=True) + if mask.any(): + out = out + mask * self.experts[expert_idx](x) + return out + + +class TinyBlock(nn.Module): + def __init__(self, d_model: int, n_heads: int, d_ff: int, mlp: nn.Module): + super().__init__() + self.norm1 = TinyRMSNorm(d_model) + self.attn = TinyAttention(d_model, n_heads) + self.norm2 = TinyRMSNorm(d_model) + self.mlp = mlp + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x), cos, sin) + x = x + self.mlp(self.norm2(x)) + return x + + +class TinyPretrainModel(nn.Module): + """Minimal decoder-only model matching the adapter's expected paths: + embed / blocks / norm_f / lm_head. `moe_layer_indices` (possibly empty) + controls which blocks get a `TinyMoE` mlp instead of a dense one. + """ + + def __init__( + self, + d_model: int = 32, + n_heads: int = 4, + n_layers: int = 2, + d_ff: int = 64, + vocab_size: int = 256, + n_experts: int = 4, + top_k: int = 2, + moe_layer_indices: frozenset[int] = frozenset(), + tie_embeddings: bool = True, + ): + super().__init__() + self.embed = nn.Embedding(vocab_size, d_model) + self.blocks = nn.ModuleList( + [ + TinyBlock( + d_model, + n_heads, + d_ff, + mlp=( + TinyMoE(d_model, d_ff, n_experts, top_k) + if i in moe_layer_indices + else TinyDenseMLP(d_model, d_ff) + ), + ) + for i in range(n_layers) + ] + ) + self.norm_f = TinyRMSNorm(d_model) + self.lm_head = nn.Linear(d_model, vocab_size, bias=False) + if tie_embeddings: + self.lm_head.weight = self.embed.weight + self.d_head = d_model // n_heads + + def _forward_impl(self, tokens: torch.Tensor) -> dict: + b, s = tokens.shape + x = self.embed(tokens) + # Dummy cos/sin (no real RoPE math needed for these structural + # tests) but real tensors, passed positionally into every block -- + # matches the real target architecture's convention. + cos = torch.ones(s, self.d_head // 2) + sin = torch.zeros(s, self.d_head // 2) + for block in self.blocks: + x = block(x, cos, sin) + x = self.norm_f(x) + return {"logits": self.lm_head(x)} + + def forward(self, tokens: torch.Tensor, **kwargs) -> dict: + # **kwargs declared so pytest fixtures can pass through + # run_with_cache's force-injected kwargs -- the actual kwarg- + # forwarding contracts (strict vs. **kwargs-accepting source + # forwards) are exercised directly against ForwardStrict/ + # ForwardVarKwargs below instead. + return self._forward_impl(tokens) + + +class ForwardStrict(nn.Module): + """Declares no **kwargs catch-all -- exercises the "must filter" + branch of PretrainModelContainer's kwarg handling.""" + + def __init__(self): + super().__init__() + self.seen_kwargs: dict = {} + + def forward(self, tokens: torch.Tensor, targets: torch.Tensor | None = None) -> dict: + self.seen_kwargs = {"targets": targets} + return {"logits": tokens.float()} + + +class ForwardVarKwargs(nn.Module): + """Declares **kwargs -- exercises the "pass everything through" branch + of PretrainModelContainer's kwarg handling.""" + + def __init__(self): + super().__init__() + self.seen_kwargs: dict = {} + + def forward(self, tokens: torch.Tensor, **kwargs) -> dict: + self.seen_kwargs = kwargs + return {"logits": tokens.float()} + + +class MalformedMLP(nn.Module): + """Has neither a dense MLP's (gate, up, down) nor an MoE layer's + (router, experts) attributes -- used to test the adapter's failure + path.""" + + def __init__(self, d_model: int): + super().__init__() + self.some_other_layer = nn.Linear(d_model, d_model) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.some_other_layer(x) + + +def make_cfg(d_model: int = 32, n_layers: int = 2) -> TransformerBridgeConfig: + return TransformerBridgeConfig( + d_model=d_model, + d_head=d_model // 4, + n_layers=n_layers, + n_ctx=128, + n_heads=4, + d_vocab=256, + d_mlp=d_model * 2, + architecture=ARCHITECTURE_NAME, + ) diff --git a/tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py b/tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py new file mode 100644 index 000000000..d2fa03f24 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_pretrain_adapter.py @@ -0,0 +1,453 @@ +"""Unit tests for PretrainArchitectureAdapter and DenseOrMoEFeedForwardBridge +-- component mapping, structural dispatch, and bridge-construction-level +behavior. Container, builder, and wrapped-model lifecycle behavior (kwarg +filtering, output-contract normalization, hook registration for the hidden +MLP delegate, train/eval, dtype/device) lives in +test_pretrain_model_container.py. + +Fully self-contained (see _pretrain_mocks.py): tiny mock `nn.Module`s, no +external package dependency, no network access. +""" + +from __future__ import annotations + +import pytest +import torch + +from transformer_lens.model_bridge.generalized_components import ( + DelegatedAttentionBlockBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, +) +from transformer_lens.model_bridge.supported_architectures.pretrain import ( + ARCHITECTURE_NAME, + DenseOrMoEFeedForwardBridge, + NativeForwardAttentionBridge, + PretrainArchitectureAdapter, + PretrainModelContainer, + build_pretrain_bridge, +) + +from ._pretrain_mocks import ( + MalformedMLP, + TinyDenseMLP, + TinyMoE, + TinyPretrainModel, + make_cfg, +) + +_make_cfg = make_cfg # local alias, matches call sites below + + +class TestPretrainAdapterConstruction: + def test_adapter_sets_required_config_flags(self) -> None: + adapter = PretrainArchitectureAdapter(_make_cfg()) + assert adapter.cfg.normalization_type == "RMS" + assert adapter.cfg.positional_embedding_type == "rotary" + assert adapter.cfg.final_rms is True + assert adapter.cfg.gated_mlp is True + assert adapter.cfg.attn_only is False + + def test_adapter_mutates_the_passed_in_cfg_object_in_place(self) -> None: + """Documented, intentional behavior (matches nanogpt.py's adapter + convention) -- not an accidental side effect. This test guards the + specific claim: the object passed in is the object mutated, not a + copy.""" + cfg = _make_cfg() + adapter = PretrainArchitectureAdapter(cfg) + assert adapter.cfg is cfg + assert cfg.normalization_type == "RMS" + + def test_component_mapping_has_expected_top_level_keys(self) -> None: + adapter = PretrainArchitectureAdapter(_make_cfg()) + mapping = adapter.get_component_mapping() + assert set(mapping.keys()) == {"embed", "blocks", "ln_final", "unembed"} + + def test_mlp_mapping_always_uses_the_dispatcher(self) -> None: + """Unconditional now -- no cfg.num_experts branch to depend on.""" + adapter = PretrainArchitectureAdapter(_make_cfg()) + mlp_component = adapter.component_mapping["blocks"].submodules["mlp"] + assert isinstance(mlp_component, DenseOrMoEFeedForwardBridge) + + +class TestNativeForwardAttentionBridge: + def test_opaque_attention_bridge_does_not_advertise_per_head_aliases(self) -> None: + """The opaque attention wrap has no Q/K/V/O submodules, so it must + not advertise AttentionBridge's class-level hook_aliases/ + property_aliases (which assume those submodules exist) or the + split-QKV-fork machinery -- see NativeForwardAttentionBridge's + docstring.""" + adapter = PretrainArchitectureAdapter(_make_cfg()) + attn = adapter.component_mapping["blocks"].submodules["attn"] + + assert isinstance(attn, NativeForwardAttentionBridge) + assert attn.hook_aliases == {} + assert attn.property_aliases == {} + assert attn.supports_split_qkv_fork is False + + # Also check the actual constructed runtime bridge, not just the + # adapter's own component_mapping -- proves bridge construction + # didn't reintroduce the aliases along the way. + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=1, d_ff=64, vocab_size=256) + bridge = build_pretrain_bridge(model, cfg) + assert bridge.blocks[0].attn.hook_aliases == {} + + def test_block_level_attention_output_still_reaches_hook_out(self) -> None: + """Clearing the per-head aliases must not also silence the block's + own attn.hook_out -- that's the actual hook point this adapter's + block-level hook coverage relies on.""" + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=1, d_ff=64, vocab_size=256) + bridge = build_pretrain_bridge(model, cfg) + tokens = torch.randint(0, 256, (1, 5)) + _, cache = bridge.run_with_cache(tokens) + assert "blocks.0.attn.hook_out" in cache + + def test_blocks_use_delegated_attention_block_bridge(self) -> None: + """The adapter reuses DelegatedAttentionBlockBridge rather than + plain BlockBridge, since NativeForwardAttentionBridge's + supports_split_qkv_fork = False means the block-level + hook_attn_in/hook_q_input/hook_k_input/hook_v_input aliases would + otherwise dangle (point at HookPoints that are never created).""" + adapter = PretrainArchitectureAdapter(_make_cfg()) + assert isinstance(adapter.component_mapping["blocks"], DelegatedAttentionBlockBridge) + + def test_split_qkv_hook_names_are_absent_but_attention_output_remains_available( + self, + ) -> None: + """hook_attn_in/hook_q_input/hook_k_input/hook_v_input must not + appear as resolvable hook names anywhere in the built bridge -- + DelegatedAttentionBlockBridge strips the block-level aliases, and + NativeForwardAttentionBridge never creates the underlying + attention-level HookPoints in the first place. hook_attn_out is + untouched by either change and must still resolve, both as the + block-level alias and as the attention component's own hook_out.""" + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=1, d_ff=64, vocab_size=256) + bridge = build_pretrain_bridge(model, cfg) + hook_dict = bridge.hook_dict + for dangling in ( + "blocks.0.hook_attn_in", + "blocks.0.hook_q_input", + "blocks.0.hook_k_input", + "blocks.0.hook_v_input", + "blocks.0.attn.hook_attn_in", + "blocks.0.attn.hook_q_input", + "blocks.0.attn.hook_k_input", + "blocks.0.attn.hook_v_input", + ): + assert dangling not in hook_dict, f"{dangling} should not resolve" + assert "blocks.0.hook_attn_out" in hook_dict + assert "blocks.0.attn.hook_out" in hook_dict + + +class TestDenseOrMoEFeedForwardBridgeDispatch: + def test_dispatches_to_moe_delegate_for_moe_module(self) -> None: + cfg = _make_cfg() + bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg) + moe_module = TinyMoE(d_model=32, d_ff=64, n_experts=4, top_k=2) + bridge.set_original_component(moe_module) + + from transformer_lens.model_bridge.generalized_components import MoEBridge + + assert isinstance(bridge._delegate, MoEBridge) + + def test_dispatches_to_dense_delegate_for_dense_module(self) -> None: + cfg = _make_cfg() + bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg) + dense_module = TinyDenseMLP(d_model=32, d_ff=64) + bridge.set_original_component(dense_module) + + from transformer_lens.model_bridge.generalized_components import GatedMLPBridge + + assert isinstance(bridge._delegate, GatedMLPBridge) + + def test_raises_clear_error_for_malformed_module(self) -> None: + cfg = _make_cfg() + bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg) + malformed = MalformedMLP(d_model=32) + + with pytest.raises(ValueError, match="doesn't know how to wrap it"): + bridge.set_original_component(malformed) + + +class TestDenseOrMoEFeedForwardBridgeDispatchValidatesTypes: + """hasattr alone would let a module win a branch by attribute-name + coincidence even if the attributes are the wrong type. These tests + exercise the basic type checks added to set_original_component so + such a mismatch is caught here, with a clear message naming the + actual field and type, rather than failing later inside MoEBridge/ + GatedMLPBridge or during forward.""" + + def test_router_of_wrong_type_raises_type_error(self) -> None: + class BadRouterMoE(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.router = "not_a_module" # right name, wrong type + self.experts = torch.nn.ModuleList( + [TinyDenseMLP(d_model=32, d_ff=64) for _ in range(4)] + ) + + cfg = _make_cfg() + bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg) + with pytest.raises(TypeError, match="router"): + bridge.set_original_component(BadRouterMoE()) + + def test_experts_of_wrong_type_raises_type_error(self) -> None: + class BadExpertsMoE(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.router = torch.nn.Linear(32, 4) + self.experts = "not_a_module_collection" # right name, wrong type + + cfg = _make_cfg() + bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg) + with pytest.raises(TypeError, match="experts"): + bridge.set_original_component(BadExpertsMoE()) + + def test_plain_list_of_experts_is_rejected(self) -> None: + """Modules held in an ordinary Python list are not registered as + children -- their parameters would silently drop out of + parameters()/state_dict()/.to(...)/train()/eval(), contradicting + this adapter's lifecycle guarantees. Rejected even though every + element is itself a valid nn.Module.""" + + class ListExpertsMoE(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.router = torch.nn.Linear(32, 4) + self.experts = [TinyDenseMLP(d_model=32, d_ff=64) for _ in range(4)] + + cfg = _make_cfg() + bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg) + with pytest.raises(TypeError, match="registered module collection"): + bridge.set_original_component(ListExpertsMoE()) + + def test_gate_of_wrong_type_raises_type_error(self) -> None: + class BadGateDense(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate = "not_a_module" # right name, wrong type + self.up = torch.nn.Linear(32, 64, bias=False) + self.down = torch.nn.Linear(64, 32, bias=False) + + cfg = _make_cfg() + bridge = DenseOrMoEFeedForwardBridge(name="mlp", config=cfg) + with pytest.raises(TypeError, match="gate"): + bridge.set_original_component(BadGateDense()) + + +class TestDenseOrMoEFeedForwardBridgeTupleOutput: + """Exercises the dispatcher's own tuple-output handling directly, by + patching an already-wired delegate's forward -- independent of + whether GatedMLPBridge/MoEBridge happen to return a tuple today. The + dispatcher explicitly supports and validates this shape, so it's + tested as its own contract rather than assumed from what the real + delegates currently do.""" + + def _dispatcher_with_patched_delegate(self, fake_forward): + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + n_experts=4, + top_k=2, + moe_layer_indices=frozenset({0}), + ) + bridge = build_pretrain_bridge(model, cfg) + dispatcher = bridge.blocks[0].mlp + dispatcher._delegate.forward = fake_forward + return dispatcher + + def test_tensor_first_tuple_passes_through_with_hook_applied(self) -> None: + dispatcher = self._dispatcher_with_patched_delegate( + lambda *a, **kw: (torch.zeros(1, 3, 16), "aux") + ) + output = dispatcher(torch.ones(1, 3, 16)) + assert isinstance(output, tuple) + assert output[1] == "aux" + torch.testing.assert_close(output[0], torch.zeros(1, 3, 16)) + + def test_empty_tuple_raises_clear_type_error(self) -> None: + dispatcher = self._dispatcher_with_patched_delegate(lambda *a, **kw: ()) + with pytest.raises(TypeError, match="torch.Tensor"): + dispatcher(torch.ones(1, 3, 16)) + + def test_non_tensor_first_element_raises_clear_type_error(self) -> None: + dispatcher = self._dispatcher_with_patched_delegate( + lambda *a, **kw: ("not_a_tensor", torch.zeros(1)) + ) + with pytest.raises(TypeError, match="torch.Tensor"): + dispatcher(torch.ones(1, 3, 16)) + + +class TestPretrainAdapterBridgeConstruction: + @pytest.fixture + def dense_bridge(self): + cfg = _make_cfg(n_layers=2) + model = TinyPretrainModel(d_model=32, n_heads=4, n_layers=2, d_ff=64, vocab_size=256) + return build_pretrain_bridge(model, cfg) + + @pytest.fixture + def moe_bridge(self): + cfg = _make_cfg(n_layers=2) + model = TinyPretrainModel( + d_model=32, + n_heads=4, + n_layers=2, + d_ff=64, + vocab_size=256, + n_experts=4, + top_k=2, + moe_layer_indices=frozenset({1}), + ) + return build_pretrain_bridge(model, cfg) + + def test_bridge_has_correct_block_count(self, dense_bridge) -> None: + assert len(dense_bridge.blocks) == 2 + + def test_bridge_has_embed_unembed_and_final_norm(self, dense_bridge) -> None: + assert hasattr(dense_bridge, "embed") + assert hasattr(dense_bridge, "unembed") + assert hasattr(dense_bridge, "ln_final") + + def test_forward_returns_logits_of_expected_shape(self, dense_bridge) -> None: + tokens = torch.randint(0, 256, (1, 5)) + with torch.no_grad(): + output = dense_bridge(tokens) + assert output.shape == (1, 5, 256) + + def test_run_with_cache_exposes_resid_hooks(self, dense_bridge) -> None: + tokens = torch.randint(0, 256, (1, 5)) + _, cache = dense_bridge.run_with_cache(tokens) + assert any("resid_pre" in k for k in cache.keys()) + assert any("resid_post" in k for k in cache.keys()) + + def test_no_per_head_attention_hooks(self, dense_bridge) -> None: + """This adapter deliberately does not expose hook_q/hook_k/hook_v -- + see PretrainArchitectureAdapter's class docstring.""" + tokens = torch.randint(0, 256, (1, 5)) + _, cache = dense_bridge.run_with_cache(tokens) + assert not any("hook_q" in k for k in cache.keys()) + assert not any("hook_pattern" in k for k in cache.keys()) + + def test_moe_bridge_forward_runs(self, moe_bridge) -> None: + tokens = torch.randint(0, 256, (1, 5)) + with torch.no_grad(): + output = moe_bridge(tokens) + assert output.shape == (1, 5, 256) + assert not torch.isnan(output).any() + + def test_moe_block_uses_dispatcher(self, moe_bridge) -> None: + assert isinstance(moe_bridge.blocks[1].mlp, DenseOrMoEFeedForwardBridge) + + def test_moe_layer_hooks_are_registered_at_the_dispatcher_path(self, moe_bridge) -> None: + """Verifies actual cache keys, not just isinstance -- the claim in + DenseOrMoEFeedForwardBridge's docstring that hooks fire at + blocks.{i}.mlp.hook_in/hook_out regardless of dense-vs-MoE.""" + tokens = torch.randint(0, 256, (1, 5)) + _, cache = moe_bridge.run_with_cache(tokens) + assert "blocks.1.mlp.hook_in" in cache + assert "blocks.1.mlp.hook_out" in cache + # Layer 0 is dense in this fixture -- same path convention applies. + assert "blocks.0.mlp.hook_in" in cache + assert "blocks.0.mlp.hook_out" in cache + + +class TestMoEConfigFieldIndependence: + def test_bridge_behaves_identically_with_moe_config_fields_populated(self) -> None: + """MoEBridge itself (transformer_lens/model_bridge/generalized_ + components/moe.py) never reads cfg.num_experts or + cfg.experts_per_token. Populating them with realistic values must + not change behavior. + Uses an actual MoE layer (moe_layer_indices={0}) -- both models + being dense would let this pass without ever constructing + MoEBridge, so the dispatch type is asserted directly too.""" + from transformer_lens.model_bridge.generalized_components import MoEBridge + + torch.manual_seed(0) + model_a = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + n_experts=4, + top_k=2, + moe_layer_indices=frozenset({0}), + ) + cfg_without = _make_cfg(n_layers=1) + bridge_without = build_pretrain_bridge(model_a, cfg_without) + + torch.manual_seed(0) + model_b = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + n_experts=4, + top_k=2, + moe_layer_indices=frozenset({0}), + ) + cfg_with = _make_cfg(n_layers=1) + cfg_with.num_experts = 4 + cfg_with.experts_per_token = 2 + bridge_with = build_pretrain_bridge(model_b, cfg_with) + + assert isinstance(bridge_without.blocks[0].mlp._delegate, MoEBridge) + assert isinstance(bridge_with.blocks[0].mlp._delegate, MoEBridge) + + tokens = torch.randint(0, 64, (1, 3)) + with torch.no_grad(): + out_without = bridge_without(tokens) + out_with = bridge_with(tokens) + torch.testing.assert_close(out_without, out_with, atol=0, rtol=0) + + +class TestPretrainAdapterTiedEmbeddings: + def test_tied_embedding_bridge_construction_succeeds(self) -> None: + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel( + d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64, tie_embeddings=True + ) + bridge = build_pretrain_bridge(model, cfg) + tokens = torch.randint(0, 64, (1, 3)) + with torch.no_grad(): + output = bridge(tokens) + assert output.shape == (1, 3, 64) + + def test_untied_embedding_bridge_construction_succeeds(self) -> None: + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel( + d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64, tie_embeddings=False + ) + bridge = build_pretrain_bridge(model, cfg) + tokens = torch.randint(0, 64, (1, 3)) + with torch.no_grad(): + output = bridge(tokens) + assert output.shape == (1, 3, 64) + + +class TestBuildBridgeFromModuleDirectlyStillWorks: + """Advanced/internal path -- build_pretrain_bridge is the intended + public entry point, but direct build_bridge_from_module use (with the + container applied manually) must keep working too.""" + + def test_direct_build_bridge_from_module_with_manual_container(self) -> None: + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge = build_bridge_from_module( + PretrainModelContainer(model), + architecture=ARCHITECTURE_NAME, + tl_config=cfg, + ) + tokens = torch.randint(0, 64, (1, 3)) + with torch.no_grad(): + output = bridge(tokens) + assert output.shape == (1, 3, 64) diff --git a/tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py b/tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py new file mode 100644 index 000000000..43c40901e --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_pretrain_model_container.py @@ -0,0 +1,537 @@ +"""Container, builder, and wrapped-model lifecycle behavior for +PretrainModelContainer and build_pretrain_bridge -- kwarg filtering, +output-contract normalization, train/eval mode propagation, and +dtype/tied-weight preservation. Hook registration tests for the hidden +MLP delegate +(implemented by DenseOrMoEFeedForwardBridge) live here too, since what +they verify -- real parameters staying reachable through +`original_model.inner` -- is a lifecycle concern, not a mapping one. +Adapter-mapping-level behavior lives in test_pretrain_adapter.py. + +Fully self-contained (see _pretrain_mocks.py): tiny mock `nn.Module`s, no +external package dependency, no network access. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from transformer_lens.model_bridge.supported_architectures.pretrain import ( + PretrainModelContainer, + _LogitsAttrDict, + build_pretrain_bridge, +) + +from ._pretrain_mocks import ( + ForwardStrict, + ForwardVarKwargs, + TinyPretrainModel, + make_cfg, +) + +_make_cfg = make_cfg # local alias, matches call sites below + + +class TestPretrainModelContainerKwargFiltering: + def test_output_attentions_is_stripped_for_strict_signature(self) -> None: + """The one TransformerBridge-injected kwarg (run_with_cache + forces output_attentions=True) must not reach a source forward with + no **kwargs catch-all.""" + model = ForwardStrict() + container = PretrainModelContainer(model) + container(torch.tensor([[1, 2, 3]]), targets=None, output_attentions=True) + assert model.seen_kwargs == {"targets": None} + + def test_output_attentions_is_stripped_even_for_var_kwargs_signature(self) -> None: + """Stripped unconditionally, regardless of whether the source would + have accepted it anyway -- keeps behavior uniform across source + forward signatures.""" + model = ForwardVarKwargs() + container = PretrainModelContainer(model) + container(torch.tensor([[1, 2, 3]]), output_attentions=True, some_other_kwarg=42) + assert model.seen_kwargs == {"some_other_kwarg": 42} + + def test_unrelated_kwarg_typo_is_not_silently_swallowed(self) -> None: + """A genuine caller mistake (e.g. `target=` instead of `targets=`) + must raise, not be silently discarded alongside the one kwarg this + container is actually responsible for stripping.""" + model = ForwardStrict() + container = PretrainModelContainer(model) + with pytest.raises(TypeError): + container(torch.tensor([[1, 2, 3]]), target=torch.tensor([1])) + + +class TestPretrainModelContainerHookRegistration: + """The delegate is hidden from nn.Module registration (see + DenseOrMoEFeedForwardBridge), so blocks.{i}.mlp.hook_in/out are the + only publicly registered MLP hook points. Forward, gradients, dtype + conversion, and state_dict() still reach the hidden delegate's real + weights through the raw wrapped model.""" + + def _build(self): + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + return build_pretrain_bridge(model, cfg) + + def test_only_the_dispatcher_hook_out_is_registered(self) -> None: + bridge = self._build() + tokens = torch.randint(0, 64, (1, 4)) + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens) + assert "blocks.0.mlp.hook_out" in cache + assert "blocks.0.mlp._delegate.hook_out" not in cache + + def test_broad_hook_selector_applies_intervention_exactly_once(self) -> None: + bridge = self._build() + tokens = torch.randint(0, 64, (1, 4)) + calls: list[str] = [] + + def recording_add_one(value: torch.Tensor, hook) -> torch.Tensor: + calls.append(hook.name) + return value + 1.0 + + def broad_filter(name: str) -> bool: + return "mlp" in name and "hook_out" in name + + with torch.no_grad(): + bridge.run_with_hooks(tokens, fwd_hooks=[(broad_filter, recording_add_one)]) + + assert calls == ["blocks.0.mlp.hook_out"] + + def test_gradients_still_reach_the_hidden_delegates_weights(self) -> None: + bridge = self._build() + tokens = torch.randint(0, 64, (1, 4)) + bridge(tokens).sum().backward() + mlp = bridge.original_model.inner.blocks[0].mlp + assert mlp.gate.weight.grad is not None + assert mlp.gate.weight.grad.abs().sum().item() > 0 + + def test_bridge_to_dtype_reaches_the_hidden_delegates_weights(self) -> None: + bridge = self._build() + bridge.to(dtype=torch.float64) + + assert bridge.original_model.inner.embed.weight.dtype == torch.float64 + assert bridge.original_model.inner.blocks[0].mlp.gate.weight.dtype == torch.float64 + + tokens = torch.randint(0, 64, (1, 4)) + with torch.no_grad(): + output = bridge(tokens) + assert output.dtype == torch.float64 + + def test_state_dict_includes_the_hidden_delegates_weights(self) -> None: + bridge = self._build() + source_weight = bridge.original_model.inner.blocks[0].mlp.gate.weight + + state = bridge.state_dict() + matching = [v for k, v in state.items() if k.endswith("blocks.0.mlp.gate.weight")] + + assert len(matching) == 1 + torch.testing.assert_close(matching[0], source_weight) + + def test_repeated_hook_lifecycle_still_works(self) -> None: + """Guards against a future assumption that every internally-used + GeneralizedComponent must be registered -- the dispatcher's own + hooks are, and that's what reset_hooks()/run_with_cache rely on.""" + bridge = self._build() + tokens = torch.randint(0, 64, (1, 4)) + with torch.no_grad(): + bridge.run_with_cache(tokens) + bridge.reset_hooks() + _, cache = bridge.run_with_cache(tokens) + assert "blocks.0.mlp.hook_out" in cache + + +class TestTrainEvalModePropagation: + """bridge.train()/.eval() propagate to the wrapped model: + build_pretrain_bridge reassigns the returned bridge's class to a small + generated TransformerBridge subclass (see + _get_mode_propagating_bridge_class) whose overridden train() also sets + mode on original_model; inherited eval() calls train(False), so it + reaches the override too without needing its own override. This + closes a gap in TransformerBridge's own train()/eval(), which only + walk component_mapping and never reach original_model on their own.""" + + def test_bridge_eval_propagates_to_wrapped_model(self) -> None: + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge = build_pretrain_bridge(model, cfg) + + assert model.training is True + bridge.eval() + assert model.training is False + + def test_bridge_train_propagates_to_wrapped_model(self) -> None: + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge = build_pretrain_bridge(model, cfg) + + model.eval() + assert model.training is False + bridge.train() + assert model.training is True + + def test_caller_can_still_set_mode_directly_on_their_own_model_reference(self) -> None: + """Setting mode via the raw model reference still works and stays + in sync -- both paths write to the same underlying module.""" + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + build_pretrain_bridge(model, cfg) + + model.eval() + assert model.training is False + model.train() + assert model.training is True + + def test_train_and_eval_return_the_bridge_itself(self) -> None: + """nn.Module convention: train()/eval() return self, so callers can + chain (`model.train().to(device)`, etc). The generated subclass + must preserve this, not just the mode-propagation side effect.""" + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge = build_pretrain_bridge(model, cfg) + + assert bridge.train(False) is bridge + assert bridge.train(True) is bridge + assert bridge.eval() is bridge + + def test_bridge_can_be_reconstructed_from_supported_state(self) -> None: + """The adapter's actual persistence contract is: rebuild the + bridge from the source model class + cfg + the source model's + *raw* state_dict, then confirm behavior (outputs, train/eval + propagation) survives the round trip. + + The state_dict must be snapshotted BEFORE `build_pretrain_bridge` + is called. Building the bridge wraps the model's submodules in + place (each one gains an `_original_component` wrapper for + hooking), so `model.state_dict()` taken *after* wrapping has + keys like `embed._original_component.weight` instead of + `embed.weight` -- a fresh, unwrapped `TinyPretrainModel` can't + load that. Capturing the state_dict first, then wrapping a + second fresh model and loading the pre-wrap snapshot into it, + avoids that trap. + + Only the source model's state_dict is used -- not + `bridge.state_dict()` -- since the bridge doesn't own an + independent set of learned weights; it exposes the wrapped + source model's parameters under bridge-shaped names. + + This is deliberately weaker than whole-object `pickle.dumps`, + which would also require every nested `TransformerBridge` + implementation detail (e.g. `AttentionBridge`'s hook + conversions) to be picklable -- a guarantee unrelated to this + adapter and not part of what it needs to promise. If a stronger + whole-object-pickle guarantee is later required, it belongs in a + `TransformerBridge`-level test, not here. + """ + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + + # Snapshot BEFORE wrapping -- see docstring. + raw_model_state = {key: value.detach().clone() for key, value in model.state_dict().items()} + + bridge = build_pretrain_bridge(model, cfg) + bridge.eval() + + restored_model = TinyPretrainModel( + d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64 + ) + restored_model.load_state_dict(raw_model_state) + + restored_bridge = build_pretrain_bridge(restored_model, cfg) + restored_bridge.eval() + + tokens = torch.randint(0, 64, (1, 4)) + with torch.no_grad(): + expected = bridge(tokens) + actual = restored_bridge(tokens) + + expected_logits = expected.logits if hasattr(expected, "logits") else expected + actual_logits = actual.logits if hasattr(actual, "logits") else actual + torch.testing.assert_close(actual_logits, expected_logits, atol=0, rtol=0) + + assert restored_bridge.training is False + assert restored_model.training is False + + restored_bridge.train() + assert restored_model.training is True + + def test_repeated_construction_reuses_the_same_generated_class(self) -> None: + """Guards the cache's stated purpose: build_pretrain_bridge should + not generate a new mode-propagating subclass on every call -- two + independently-built bridges from the same base bridge class must + share one generated class.""" + cfg_a = _make_cfg(n_layers=1) + model_a = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge_a = build_pretrain_bridge(model_a, cfg_a) + + cfg_b = _make_cfg(n_layers=1) + model_b = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge_b = build_pretrain_bridge(model_b, cfg_b) + + assert type(bridge_a) is type(bridge_b) + + +class TestDtypeAndTiedWeightPreservation: + def test_non_default_dtype_and_tied_weights_survive_construction(self) -> None: + torch.manual_seed(0) + model = TinyPretrainModel( + d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64, tie_embeddings=True + ).to(torch.float64) + assert model.lm_head.weight is model.embed.weight + + cfg = _make_cfg(n_layers=1) + bridge = build_pretrain_bridge(model, cfg) + + assert next(model.parameters()).dtype == torch.float64 + assert model.lm_head.weight is model.embed.weight + + tokens = torch.randint(0, 64, (1, 3)) + with torch.no_grad(): + output = bridge(tokens) + assert output.dtype == torch.float64 + + +class TestBuildPretrainBridgeErgonomics: + def test_explicit_model_name_is_forwarded(self) -> None: + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge = build_pretrain_bridge(model, cfg, model_name="my-custom-name") + assert bridge.cfg.model_name == "my-custom-name" + + def test_omitting_optional_kwargs_still_works(self) -> None: + """The default (no device/dtype/model_name passed) path, exercised + everywhere else in this file, must keep working unchanged.""" + cfg = _make_cfg(n_layers=1) + model = TinyPretrainModel(d_model=16, n_heads=2, n_layers=1, d_ff=32, vocab_size=64) + bridge = build_pretrain_bridge(model, cfg) + tokens = torch.randint(0, 64, (1, 3)) + with torch.no_grad(): + output = bridge(tokens) + assert output.shape == (1, 3, 64) + + +class TestPretrainModelContainerOutputContract: + def test_container_normalizes_dict_output_to_logits_attr_dict(self) -> None: + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + ) + container = PretrainModelContainer(model) + tokens = torch.randint(0, 64, (1, 3)) + + with torch.no_grad(): + output = container(tokens) + + assert isinstance(output, _LogitsAttrDict) + assert output.logits is output["logits"] + assert output.logits.shape == (1, 3, 64) + + def test_container_returns_a_fresh_wrapper_on_each_call(self) -> None: + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + ) + container = PretrainModelContainer(model) + tokens = torch.randint(0, 64, (1, 3)) + + with torch.no_grad(): + out1 = container(tokens) + out2 = container(tokens) + + assert isinstance(out1, _LogitsAttrDict) + assert isinstance(out2, _LogitsAttrDict) + assert out1 is not out2 + assert out1["logits"] is not out2["logits"] + torch.testing.assert_close(out1["logits"], out2["logits"]) + + def test_container_does_not_reuse_or_mutate_output_mapping(self) -> None: + model = TinyPretrainModel( + d_model=16, + n_heads=2, + n_layers=1, + d_ff=32, + vocab_size=64, + ).eval() + container = PretrainModelContainer(model) + tokens = torch.randint(0, 64, (1, 3)) + + with torch.no_grad(): + out1 = container(tokens) + + out1["sentinel"] = True + + with torch.no_grad(): + out2 = container(tokens) + + assert out1 is not out2 + assert "sentinel" not in out2 + torch.testing.assert_close(out1["logits"], out2["logits"]) + + def test_dict_missing_logits_key_raises_value_error(self) -> None: + class ForwardWrongKey(nn.Module): + def forward(self, tokens: torch.Tensor) -> dict: + return {"scores": tokens.float()} + + container = PretrainModelContainer(ForwardWrongKey()) + with pytest.raises(ValueError, match="'logits'"): + container(torch.tensor([[1, 2, 3]])) + + def test_dict_with_heterogeneous_keys_still_raises_value_error(self) -> None: + """sorted(output.keys()) would raise TypeError on non-comparable + mixed key types (e.g. str and int); the error message must use + list() instead, so an invalid model output produces the intended + ValueError, not an unrelated sorting error.""" + + class ForwardHeterogeneousKeys(nn.Module): + def forward(self, tokens: torch.Tensor) -> dict: + return {"scores": tokens.float(), 1: "metadata"} + + container = PretrainModelContainer(ForwardHeterogeneousKeys()) + with pytest.raises(ValueError, match="'logits'"): + container(torch.tensor([[1, 2, 3]])) + + def test_dict_with_non_tensor_logits_raises_type_error(self) -> None: + class ForwardBadLogitsType(nn.Module): + def forward(self, tokens: torch.Tensor) -> dict: + return {"logits": "not a tensor"} + + container = PretrainModelContainer(ForwardBadLogitsType()) + with pytest.raises(TypeError, match="torch.Tensor"): + container(torch.tensor([[1, 2, 3]])) + + def test_already_logits_attr_dict_passes_through_by_identity(self) -> None: + """An already-wrapped output must not be re-wrapped.""" + already_wrapped = _LogitsAttrDict({"logits": torch.zeros(1)}) + + class ForwardAlreadyWrapped(nn.Module): + def forward(self, tokens: torch.Tensor) -> _LogitsAttrDict: + return already_wrapped + + container = PretrainModelContainer(ForwardAlreadyWrapped()) + output = container(torch.tensor([[1, 2, 3]])) + assert output is already_wrapped + + def test_malformed_logits_attr_dict_raises_value_error_not_keyerror(self) -> None: + """Not a realistic target-model failure (_LogitsAttrDict is + private), but closes the contract: a malformed instance gets the + same clear ValueError an ordinary dict missing 'logits' gets, + rather than a raw KeyError leaking out.""" + malformed = _LogitsAttrDict({"scores": torch.zeros(1)}) + + class ForwardMalformedWrapped(nn.Module): + def forward(self, tokens: torch.Tensor) -> _LogitsAttrDict: + return malformed + + container = PretrainModelContainer(ForwardMalformedWrapped()) + with pytest.raises(ValueError, match="'logits'"): + container(torch.tensor([[1, 2, 3]])) + + def test_hf_style_object_with_logits_attribute_passes_through_unchanged(self) -> None: + """A source model that already satisfies hasattr(output, 'logits') + on its own (an HF-ModelOutput-like object, not a dict) needs no + normalization at all.""" + + class FakeModelOutput: + def __init__(self, logits: torch.Tensor): + self.logits = logits + + already_hf_style = FakeModelOutput(torch.zeros(1)) + + class ForwardHFStyle(nn.Module): + def forward(self, tokens: torch.Tensor) -> FakeModelOutput: + return already_hf_style + + container = PretrainModelContainer(ForwardHFStyle()) + output = container(torch.tensor([[1, 2, 3]])) + assert output is already_hf_style + + def test_property_backed_logits_is_only_evaluated_once(self) -> None: + """.logits must be read once and cached locally, not re-evaluated + on every access -- matters for a property-backed .logits, which + could otherwise do redundant (or side-effecting) work per access.""" + access_count = 0 + + class PropertyBackedOutput: + @property + def logits(self) -> torch.Tensor: + nonlocal access_count + access_count += 1 + return torch.zeros(1) + + class ForwardPropertyBacked(nn.Module): + def forward(self, tokens: torch.Tensor) -> PropertyBackedOutput: + return PropertyBackedOutput() + + container = PretrainModelContainer(ForwardPropertyBacked()) + container(torch.tensor([[1, 2, 3]])) + assert access_count == 1 + + def test_object_with_non_tensor_logits_attribute_raises_type_error(self) -> None: + class BadModelOutput: + logits = "not a tensor" + + class ForwardBadModelOutput(nn.Module): + def forward(self, tokens: torch.Tensor) -> BadModelOutput: + return BadModelOutput() + + container = PretrainModelContainer(ForwardBadModelOutput()) + with pytest.raises(TypeError, match=r"\.logits"): + container(torch.tensor([[1, 2, 3]])) + + def test_bare_tensor_output_passes_through_unchanged(self) -> None: + class ForwardBareTensor(nn.Module): + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + return tokens.float() + + container = PretrainModelContainer(ForwardBareTensor()) + tokens = torch.tensor([[1, 2, 3]]) + output = container(tokens) + assert torch.equal(output, tokens.float()) + + def test_tuple_with_tensor_first_element_passes_through_unchanged(self) -> None: + """TransformerBridge extracts output[0] as logits for tuple + returns -- no normalization needed.""" + + class ForwardTuple(nn.Module): + def forward(self, tokens: torch.Tensor) -> tuple: + return (tokens.float(), "some_aux_value") + + container = PretrainModelContainer(ForwardTuple()) + output = container(torch.tensor([[1, 2, 3]])) + assert isinstance(output, tuple) + assert output[1] == "some_aux_value" + + def test_tuple_with_non_tensor_first_element_raises_type_error(self) -> None: + class ForwardBadTuple(nn.Module): + def forward(self, tokens: torch.Tensor) -> tuple: + return ("not a tensor", tokens) + + container = PretrainModelContainer(ForwardBadTuple()) + with pytest.raises(TypeError, match="torch.Tensor"): + container(torch.tensor([[1, 2, 3]])) + + def test_list_output_raises_type_error(self) -> None: + class ForwardList(nn.Module): + def forward(self, tokens: torch.Tensor) -> list: + return [tokens.float()] + + container = PretrainModelContainer(ForwardList()) + with pytest.raises(TypeError, match="must return"): + container(torch.tensor([[1, 2, 3]])) + + def test_string_output_raises_type_error(self) -> None: + class ForwardString(nn.Module): + def forward(self, tokens: torch.Tensor) -> str: + return "wrong" + + container = PretrainModelContainer(ForwardString()) + with pytest.raises(TypeError, match="must return"): + container(torch.tensor([[1, 2, 3]])) diff --git a/tests/unit/model_bridge/test_bridge_train_mode_propagation.py b/tests/unit/model_bridge/test_bridge_train_mode_propagation.py new file mode 100644 index 000000000..4d08d0c71 --- /dev/null +++ b/tests/unit/model_bridge/test_bridge_train_mode_propagation.py @@ -0,0 +1,100 @@ +"""Regression tests for TransformerBridge train/eval mode propagation. + +`original_model` is stored outside the registered module tree, so +inherited `nn.Module.train()` does not reach it. These tests verify that +bridge mode changes propagate to the wrapped model, whose mode-dependent +layers (Dropout) rely on that flag. Architecture-independent stub bridge; +no HF Hub access, no weights. +""" +from __future__ import annotations + +import torch +import torch.nn as nn + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import LinearBridge + + +class _StubModel(nn.Module): + """Minimal source model with a mode-dependent layer (Dropout), so + train/eval propagation has an observable behavioral consequence, not + just a flag.""" + + def __init__(self) -> None: + super().__init__() + self.proj = nn.Linear(4, 4) + self.drop = nn.Dropout(0.5) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.drop(self.proj(x)) + + +class _StubAdapter(ArchitectureAdapter): + """Smallest adapter TransformerBridge will accept: a one-entry + component_mapping. The key is `stub_proj` (not `proj` or `embed`) + because TransformerBridge reserves several canonical attribute names + on itself; a colliding key fails at add_module time.""" + + def __init__(self, cfg: TransformerBridgeConfig) -> None: + super().__init__(cfg) + self.component_mapping = {"stub_proj": LinearBridge(name="proj")} + + +def _make_stub_bridge() -> tuple[TransformerBridge, _StubModel]: + cfg = TransformerBridgeConfig( + d_model=4, + d_head=2, + n_layers=1, + n_ctx=8, + n_heads=2, + d_vocab=8, + architecture="StubForTest", + ) + model = _StubModel() + bridge = TransformerBridge(model, _StubAdapter(cfg), tokenizer=None) + return bridge, model + + +class TestTrainEvalModePropagation: + """bridge.train()/.eval() must reach original_model. + + Without the TransformerBridge.train() override these fail with + original_model.training stuck at its wrap-time value (bridge.training + flips, the wrapped model's does not) -- i.e. bridge.eval() would leave + dropout active.""" + + def test_eval_propagates_to_original_model(self) -> None: + bridge, model = _make_stub_bridge() + assert model.training is True # nn.Module default at wrap time + bridge.eval() + assert bridge.training is False + assert model.training is False + + def test_train_propagates_to_original_model(self) -> None: + bridge, model = _make_stub_bridge() + bridge.eval() + assert model.training is False + bridge.train() + assert bridge.training is True + assert model.training is True + + def test_train_and_eval_return_the_bridge_itself(self) -> None: + """nn.Module convention: train()/eval() return self so callers can + chain (`bridge.eval().to(device)`, etc.).""" + bridge, _ = _make_stub_bridge() + assert bridge.train(False) is bridge + assert bridge.train(True) is bridge + assert bridge.eval() is bridge + + def test_direct_mode_on_original_model_stays_in_sync(self) -> None: + """Setting mode via the caller's own reference to the source model + still works -- both paths write the same underlying flags, so + neither clobbers the other.""" + bridge, model = _make_stub_bridge() + bridge.eval() + model.train() + assert model.training is True + bridge.eval() + assert model.training is False diff --git a/tests/unit/tools/test_model_registry.py b/tests/unit/tools/test_model_registry.py index 8fbb12c01..53627c8fa 100644 --- a/tests/unit/tools/test_model_registry.py +++ b/tests/unit/tools/test_model_registry.py @@ -807,6 +807,7 @@ class TestRegistrySyncedWithFactory: "NeelSoluOldForCausalLM", "GPT2LMHeadCustomModel", "TransformerLensNative", + "TransformerLensPretrain", # Group 2: factory-internal alias casings (HF emits the canonical name). "Gemma1ForCausalLM", # HF emits: GemmaForCausalLM "NeoForCausalLM", # HF emits: GPTNeoForCausalLM diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index bfe1b90ff..4f1e099dd 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -73,6 +73,7 @@ Phi3ArchitectureAdapter, PhiArchitectureAdapter, PhiMoEArchitectureAdapter, + PretrainArchitectureAdapter, Qwen2ArchitectureAdapter, Qwen2MoeArchitectureAdapter, Qwen3_5ArchitectureAdapter, @@ -179,6 +180,7 @@ "Zamba2ForCausalLM": Zamba2ArchitectureAdapter, "NanoGPTForCausalLM": NanogptArchitectureAdapter, "TransformerLensNative": NativeArchitectureAdapter, + "TransformerLensPretrain": PretrainArchitectureAdapter, "MinGPTForCausalLM": MingptArchitectureAdapter, "GPTNeoForCausalLM": NeoArchitectureAdapter, "GPTNeoXForCausalLM": NeoxArchitectureAdapter, diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 3c470cf72..7b2239532 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -3703,6 +3703,14 @@ def mps(self) -> "TransformerBridge": """ return self.to(torch.device("mps")) + def train(self, mode: bool = True) -> "TransformerBridge": + """Set training mode, propagating to the wrapped source model.""" + super().train(mode) + original = getattr(self, "original_model", None) + if isinstance(original, torch.nn.Module): + original.train(mode) + return self + def add_hook( self, name: Union[str, Callable[[str], bool]], diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py index 18bbdac7a..67a677ffc 100644 --- a/transformer_lens/model_bridge/generalized_components/moe.py +++ b/transformer_lens/model_bridge/generalized_components/moe.py @@ -2,6 +2,7 @@ This module contains the bridge component for Mixture of Experts layers. """ + from __future__ import annotations from typing import Any, Dict, Optional @@ -110,10 +111,24 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: kwargs = {**kwargs, "hidden_states": hooked} output = self.original_component(*args, **kwargs) if isinstance(output, tuple): + if not output: + raise TypeError( + f"{self.name}: expected a non-empty tuple whose first " + "element is a torch.Tensor from the wrapped MoE component, " + "got an empty tuple." + ) + hidden_states = output[0] + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"{self.name}: expected the first tuple element from the " + f"wrapped MoE component to be a torch.Tensor, got " + f"{type(hidden_states).__name__}." + ) if len(output) > 1: router_scores = output[1] - self.hook_router_scores(router_scores) + if isinstance(router_scores, torch.Tensor): + self.hook_router_scores(router_scores) hidden_states = self.hook_out(hidden_states) return (hidden_states,) + output[1:] else: diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 0eac330e2..cc79dd117 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -190,6 +190,9 @@ from transformer_lens.model_bridge.supported_architectures.phimoe import ( PhiMoEArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.pretrain import ( + PretrainArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.qwen import ( QwenArchitectureAdapter, ) @@ -303,6 +306,7 @@ "PhiArchitectureAdapter", "Phi3ArchitectureAdapter", "PhiMoEArchitectureAdapter", + "PretrainArchitectureAdapter", "QwenArchitectureAdapter", "Qwen2ArchitectureAdapter", "Qwen2MoeArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/pretrain.py b/transformer_lens/model_bridge/supported_architectures/pretrain.py new file mode 100644 index 000000000..20091b65f --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/pretrain.py @@ -0,0 +1,442 @@ +"""Architecture adapter for a lightweight decoder-only pretraining model. + +Maps a decoder-only transformer using RoPE, RMSNorm, gated SwiGLU MLPs, and +optional sparse mixture-of-experts feed-forward layers into +TransformerBridge, by wrapping the source module and delegating to its own +`forward` rather than translating parameters into a second implementation. + +Usage: `build_pretrain_bridge(model, cfg)` -- the public entry point. +`PretrainModelContainer` and direct `build_bridge_from_module` use are +internal/advanced details (see `PretrainModelContainer`'s docstring). + +Scope: maps a live module into TransformerBridge. Does not load +checkpoints, merge tensor-parallel shards, or depend on a training +framework. + +Required module protocol -- "lightweight decoder-only pretraining models" +describes intent, not a generality guarantee. The wrapped model must +expose: + + model.embed (embedding lookup) + model.blocks[i].norm1 (pre-attention norm) + model.blocks[i].attn (called as attn(x, ...)) + model.blocks[i].norm2 (pre-MLP norm) + model.blocks[i].mlp (gate/up/down, or router/experts) + model.norm_f (final norm) + model.lm_head (unembedding) + +`gate`/`up`/`down` and `router`/`experts` name the supported protocol. +`DenseOrMoEFeedForwardBridge` checks these structurally -- attribute +presence plus basic type (each is a module, `experts` is a registered +module collection) -- and raises clearly on a mismatch, but that is +structural validation only: it does not and cannot validate that a +module satisfying the shape actually implements matching forward +semantics. Blocks must take more than the bare hidden state (this target +passes `cos`/`sin`) -- see `PretrainModelContainer`. +""" +from __future__ import annotations + +from typing import Any + +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + AttentionBridge, + DelegatedAttentionBlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + MoEBridge, + RMSNormalizationBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + +ARCHITECTURE_NAME = "TransformerLensPretrain" + +# Reserved bridge kwargs removed before forwarding to the wrapped model. +# Only known bridge-compatibility kwargs are stripped so genuine caller +# mistakes (e.g. `target=` vs `targets=`) still raise naturally, and no +# signature introspection is needed to support an arbitrary forward. +_BRIDGE_COMPAT_KWARGS = frozenset({"output_attentions"}) + + +class DenseOrMoEFeedForwardBridge(GeneralizedComponent): + """Wraps either a dense SwiGLU MLP or a sparse MoE layer behind a + common interface. Dispatch is determined by structural inspection + (`router`/`experts` vs `gate`/`up`/`down`) rather than configuration, + so dense, MoE, and mixed architectures all use the same component + mapping. This identifies the supported protocol -- it does not + validate that a module merely sharing those attribute names actually + implements the matching forward behavior. + + `hasattr` alone would let a module with, say, both `router`/`experts` + and `gate`/`up`/`down` (or `router`/`experts` of the wrong types) win + the MoE branch by attribute-name coincidence and fail later with a + confusing error from deep inside `MoEBridge`, or not fail until + forward time. `set_original_component` therefore also checks the + basic shape of whichever protocol wins: `router`/`gate`/`up`/`down` + must themselves be modules, and `experts` must be a *registered* + module collection (`nn.ModuleList`/`nn.ModuleDict`) -- a plain Python + list/tuple of `nn.Module` experts is rejected even though every + element is itself a valid module, because modules held in an + ordinary list aren't registered as children and would silently drop + out of `parameters()`/`state_dict()`/`.to(...)`/`train()`/`eval()`, + contradicting this adapter's lifecycle guarantees. + """ + + def __init__(self, name: str, config: Any): + super().__init__(name, config=config, submodules={}) + self._delegate: GeneralizedComponent | None = None + + def set_original_component(self, component: torch.nn.Module) -> None: + super().set_original_component(component) + # This subclass's own __init__ takes `name: str` (non-optional), so + # self.name is always a str here -- but the base GeneralizedComponent + # attribute is typed `str | None`, which is all mypy sees without this + # narrowing. MoEBridge/GatedMLPBridge both require a plain `str` name. + assert self.name is not None + if hasattr(component, "router") and hasattr(component, "experts"): + if not isinstance(component.router, torch.nn.Module): + raise TypeError( + f"{type(component).__name__}.router must be an nn.Module; " + f"got {type(component.router).__name__}." + ) + if not isinstance(component.experts, (torch.nn.ModuleList, torch.nn.ModuleDict)): + raise TypeError( + f"{type(component).__name__}.experts must be a registered " + "module collection (nn.ModuleList or nn.ModuleDict); got " + f"{type(component.experts).__name__}." + ) + delegate: GeneralizedComponent = MoEBridge( + name=self.name, + config=self.config, + submodules={"gate": LinearBridge(name="router")}, + ) + elif hasattr(component, "gate") and hasattr(component, "up") and hasattr(component, "down"): + for field in ("gate", "up", "down"): + value = getattr(component, field) + if not isinstance(value, torch.nn.Module): + raise TypeError( + f"{type(component).__name__}.{field} must be an " + f"nn.Module; got {type(value).__name__}." + ) + delegate = GatedMLPBridge( + name=self.name, + config=self.config, + submodules={ + "gate": LinearBridge(name="gate"), + "in": LinearBridge(name="up"), + "out": LinearBridge(name="down"), + }, + ) + else: + raise ValueError( + f"Block.mlp is a {type(component).__name__} with neither " + "an MoE layer's (router, experts) nor a gated MLP's " + "(gate, up, down) attributes -- this adapter doesn't know " + "how to wrap it." + ) + delegate.set_original_component(component) + # Not `self._delegate = delegate`: normal registration would + # duplicate hook_in/hook_out under a nested `._delegate.` path, + # risking a broad hook selector firing twice. + # `_delegate` is an execution helper, absent from named_modules() + # -- safe since parameters/state_dict/dtype are read from the raw + # wrapped model (see PretrainModelContainer), not this tree. + object.__setattr__(self, "_delegate", delegate) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + assert self._delegate is not None, f"{self.name}: original component not set" + if args: + args = (self.hook_in(args[0]),) + args[1:] + elif "hidden_states" in kwargs: + kwargs = {**kwargs, "hidden_states": self.hook_in(kwargs["hidden_states"])} + output = self._delegate(*args, **kwargs) + + if isinstance(output, tuple): + if len(output) == 0: + raise TypeError( + "DenseOrMoEFeedForwardBridge expected a non-empty tuple " + "whose first element is a torch.Tensor" + ) + + first = output[0] + + if not isinstance(first, torch.Tensor): + raise TypeError( + "DenseOrMoEFeedForwardBridge expected the first tuple element " + f"to be a torch.Tensor, got {type(first).__name__}" + ) + + hooked_first = self.hook_out(first) + + # Preserve every auxiliary element without sending it through HookPoint. + return (hooked_first, *output[1:]) + + if not isinstance(output, torch.Tensor): + raise TypeError( + "DenseOrMoEFeedForwardBridge expected a torch.Tensor or a tuple " + f"whose first element is a torch.Tensor, got {type(output).__name__}" + ) + + return self.hook_out(output) + + +class _LogitsAttrDict(dict): + """Makes a plain dict's keys accessible as attributes (`d.logits` reads + `d["logits"]`), so a source model's plain-dict forward output satisfies + the `hasattr(output, "logits")` contract `TransformerBridge` expects. + Behaves as a plain dict everywhere else (indexing, `.get`, `in`, ...). + """ + + def __getattr__(self, key: str) -> Any: + try: + return self[key] + except KeyError as e: + raise AttributeError(key) from e + + +class PretrainModelContainer(torch.nn.Module): + """Internal detail -- `build_pretrain_bridge` applies this + automatically. Three responsibilities, all in this container's own + `forward`: + + 1. Avoids a `TransformerBridge.__getattr__` collision: a source + model's own `self.embed`/`self.blocks` clashes with identically + named component_mapping keys. Wrapping one level deeper + (`container.inner.embed`) fixes this without touching the source + model. + 2. Normalizes the return value to the `.logits` contract + `TransformerBridge` expects: a plain `"logits"` dict (the target + architecture's actual shape) is wrapped in `_LogitsAttrDict`; a bare + tensor, tensor-first tuple, or object already exposing `.logits` + passes through after validating the tensor is present and is a + tensor; anything else raises immediately with a clear message. + 3. Strips `_BRIDGE_COMPAT_KWARGS` from kwargs before calling the + wrapped model (see that constant's comment). + + Also sidesteps a `BlockBridge` convention where a bare-tensor block + output gets wrapped in a 1-tuple for "standalone hidden_states calls": + since this target's blocks take `cos`/`sin` too, that path never + triggers, so the source forward loop needs no changes to be bridged. + + `self.inner` is a regular registered submodule, so + `container.train()`/`.eval()` already recurse into it via the normal + `nn.Module` traversal -- no override needed here. The propagation gap + lives one level up, at `TransformerBridge` itself (see + `build_pretrain_bridge`), whose `.train()`/`.eval()` do not walk down + to `original_model`. + """ + + def __init__(self, model: torch.nn.Module) -> None: + super().__init__() + self.inner = model + + def forward(self, *args: Any, **kwargs: Any) -> Any: + filtered = {k: v for k, v in kwargs.items() if k not in _BRIDGE_COMPAT_KWARGS} + output = self.inner(*args, **filtered) + + # Already-normalized or already-HF-style outputs pass through + # unchanged, but only after checking .logits is actually a tensor + # -- an object merely exposing the attribute isn't enough. + if isinstance(output, _LogitsAttrDict): + if "logits" not in output: + raise ValueError( + f"{type(self.inner).__name__}.forward returned a " + "_LogitsAttrDict without a 'logits' key." + ) + if not isinstance(output["logits"], torch.Tensor): + raise TypeError( + f"{type(self.inner).__name__}.forward returned a " + f"_LogitsAttrDict with a non-tensor 'logits' value: " + f"{type(output['logits']).__name__}." + ) + return output + + # try/except, not hasattr(): hasattr() would evaluate a + # property-backed .logits once, then a separate read would + # evaluate it again. This reads it exactly once. + try: + logits = output.logits + except AttributeError: + pass + else: + if not isinstance(logits, torch.Tensor): + raise TypeError( + f"{type(self.inner).__name__}.forward returned a " + f"{type(output).__name__} with a non-tensor .logits value: " + f"{type(logits).__name__}." + ) + return output + + if isinstance(output, torch.Tensor): + return output + + # TransformerBridge extracts output[0] as logits for tuple returns. + if isinstance(output, tuple): + if output and isinstance(output[0], torch.Tensor): + return output + raise TypeError( + f"{type(self.inner).__name__}.forward returned a tuple whose " + f"first element is a {type(output[0]).__name__ if output else 'empty tuple'}, " + "not a torch.Tensor -- TransformerBridge extracts output[0] as " + "logits for tuple returns, so it must be a tensor." + ) + + # The primary case: a plain dict, normalized into _LogitsAttrDict. + if isinstance(output, dict): + if "logits" not in output: + raise ValueError( + f"{type(self.inner).__name__}.forward returned a dict with keys " + # list(), not sorted(): sorted() raises TypeError on + # heterogeneous keys, which would mask this error. + f"{list(output.keys())}, but PretrainModelContainer requires a " + "'logits' key -- without it, TransformerBridge's own " + "hasattr(output, 'logits') check would silently fail the same " + "way this container exists to prevent." + ) + if not isinstance(output["logits"], torch.Tensor): + raise TypeError( + f"{type(self.inner).__name__}.forward returned a dict whose " + f"'logits' value is a {type(output['logits']).__name__}, not a " + "torch.Tensor." + ) + return _LogitsAttrDict(output) + + raise TypeError( + f"{type(self.inner).__name__}.forward must return a torch.Tensor, a " + "tuple whose first element is a tensor, a dict containing a " + "tensor-valued 'logits' key, or an object with a tensor-valued " + f".logits attribute; got {type(output).__name__}." + ) + + +class NativeForwardAttentionBridge(AttentionBridge): + """Opaque attention bridge that delegates to the source attention. + + This adapter intentionally exposes only input/output attention hooks. + It has no mapped Q/K/V/O projection components, so the standard + per-head aliases and weight aliases do not apply. + """ + + hook_aliases = {} + property_aliases = {} + supports_split_qkv_fork = False + + +class PretrainArchitectureAdapter(ArchitectureAdapter): + """Adapter for a decoder-only transformer using RoPE, RMSNorm, gated + SwiGLU MLPs, and optional sparse MoE feed-forward layers. + + Uses an opaque `NativeForwardAttentionBridge` with no attention + projection submodules, not `JointQKVAttentionBridge`/ + `PositionEmbeddingsAttentionBridge`: those reimplement RoPE via HF's + rotate-half convention, wrong for a source model using the + adjacent-pair convention. The opaque bridge delegates unchanged to + `Attention.forward`, so RoPE runs as written -- at the cost of no + per-head hooks, only block-level + `resid_pre`/`resid_mid`/`resid_post`. + + Blocks use `DelegatedAttentionBlockBridge` rather than plain + `BlockBridge`: that existing abstraction already exists for + architectures where attention is delegated wholesale and the + split-qkv-fork block-level aliases (`hook_attn_in`/`hook_q_input`/ + `hook_k_input`/`hook_v_input`) don't apply. It complements + `NativeForwardAttentionBridge.supports_split_qkv_fork = False` (which + prevents the split-QKV-fork machinery and its associated HookPoints + from being exposed for this attention component) by also removing the + now-dangling block-level aliases that would otherwise point at them. + `hook_attn_out` is untouched by either change, since the attention + component still fires its own `hook_out` normally. + + `self.cfg` is mutated in place, not copied (matches `nanogpt.py`'s + convention) -- callers holding another reference to the same config + will see these fields change. + + Bridges built through `build_pretrain_bridge` are given a + mode-propagating subclass so `.train()`/`.eval()` reach the wrapped + source model (see that function's docstring) -- this adapter class + itself has no lifecycle behavior of its own. + """ + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + + self.cfg.normalization_type = "RMS" + self.cfg.positional_embedding_type = "rotary" + self.cfg.final_rms = True + self.cfg.gated_mlp = True + self.cfg.attn_only = False + + self.component_mapping = { + # "inner." because this adapter expects the source model to + # arrive wrapped in `PretrainModelContainer` -- see that + # class's docstring for why. + "embed": EmbeddingBridge(name="inner.embed"), + "blocks": DelegatedAttentionBlockBridge( + name="inner.blocks", + config=self.cfg, + submodules={ + "ln1": RMSNormalizationBridge(name="norm1", config=self.cfg), + "attn": NativeForwardAttentionBridge( + name="attn", + config=self.cfg, + submodules={}, # opaque wrap -- see class docstring + ), + "ln2": RMSNormalizationBridge(name="norm2", config=self.cfg), + "mlp": DenseOrMoEFeedForwardBridge(name="mlp", config=self.cfg), + }, + ), + "ln_final": RMSNormalizationBridge(name="inner.norm_f", config=self.cfg), + "unembed": UnembeddingBridge(name="inner.lm_head"), + } + + +def build_pretrain_bridge( + model: torch.nn.Module, + cfg: TransformerBridgeConfig, + *, + device: Any = None, + dtype: torch.dtype | None = None, + model_name: str | None = None, +) -> TransformerBridge: + """Public entry point: wraps `model` in `PretrainModelContainer` and + builds a `TransformerBridge` around it. Prefer this over calling + `build_bridge_from_module` directly -- the container is easy to forget. + + `device`/`dtype`/`model_name` forward to `build_bridge_from_module` + only when explicitly given. + + `bridge.train()`/`.eval()` propagate to `model` via + `TransformerBridge.train()` itself, which sets mode on + `original_model` in addition to the registered module tree + (`original_model` is deliberately not a registered submodule, so + `nn.Module.train()`'s own recursion never reaches it). This adapter + needs nothing extra for mode propagation. + + Setting mode on `model` directly still works too and stays in sync. + """ + from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, + ) + + kwargs: dict[str, Any] = {} + if device is not None: + kwargs["device"] = device + if dtype is not None: + kwargs["dtype"] = dtype + if model_name is not None: + kwargs["model_name"] = model_name + + return build_bridge_from_module( + PretrainModelContainer(model), + architecture=ARCHITECTURE_NAME, + tl_config=cfg, + **kwargs, + )