diff --git a/tests/QUARANTINES.md b/tests/QUARANTINES.md index 5a587ea38..0026b8867 100644 --- a/tests/QUARANTINES.md +++ b/tests/QUARANTINES.md @@ -46,6 +46,7 @@ Rule ([AGENTS.md §10](../AGENTS.md#10-hard-rules)): **never add `xfail` / `skip | [`unit/model_bridge/supported_architectures/test_gemma2_adapter.py`:49](unit/model_bridge/supported_architectures/test_gemma2_adapter.py) | "Network/disk fetch of tiny Gemma2" | | [`integration/model_bridge/test_bridge_integration.py`:801](integration/model_bridge/test_bridge_integration.py) | "Skip Gemma2 in CI to avoid timeout" | | [`acceptance/model_bridge/compatibility/test_hook_completeness.py`:156](acceptance/model_bridge/compatibility/test_hook_completeness.py) | "Gemma2 too large for CI" | +| [`integration/model_bridge/test_ouro_adapter.py`:35](integration/model_bridge/test_ouro_adapter.py) module-level (+ `slow` marker) | "ByteDance/Ouro-1.4B: 2.8GB download + ~11GB RAM, too large for CI" | **Un-skip:** locally with `HF_TOKEN` sourced. diff --git a/tests/integration/model_bridge/test_ouro_adapter.py b/tests/integration/model_bridge/test_ouro_adapter.py new file mode 100644 index 000000000..48f46b013 --- /dev/null +++ b/tests/integration/model_bridge/test_ouro_adapter.py @@ -0,0 +1,235 @@ +"""Integration tests for the Ouro architecture adapter (OuroForCausalLM). + +Model: ByteDance/Ouro-1.4B + - Remote-code (auto_map to modeling_ouro), so no tiny random checkpoint + exists; this loads the real 1.4B weights (~2.8GB download, ~11GB RAM for + the two fp32 copies). Gated out of CI on network/memory budget; run + locally with: + uv run pytest tests/integration/model_bridge/test_ouro_adapter.py -v -m slow + +Ouro is a looped-depth (Universal Transformer) model: the remote-code +OuroModel.forward applies the shared 24-layer stack total_ut_steps=4 times, +with model.norm after each pass. The bridge delegates the loop to HF's own +forward, so logit parity holds unchanged; the one observable difference from +a standard decoder is that each physical block's hooks fire once per UT step +(pinned by TestOuroLoopedDepth below) and a cache keeps the final step only. + +The bridge always reimplements attention in eager mode (so the score and +pattern hooks fire), so the reference HF model is loaded independently with +attn_implementation="eager" too. Comparing against an independent HF load +(never bridge.original_model, whose modules are hook-wrapped) keeps the +parity check honest. +""" + +import os +import platform + +import pytest +import torch +from transformers import AutoConfig, AutoModelForCausalLM + +from transformer_lens.model_bridge import TransformerBridge + +MODEL = "ByteDance/Ouro-1.4B" + +pytestmark = [ + pytest.mark.slow, + pytest.mark.skipif( + bool(os.getenv("CI")), + reason="ByteDance/Ouro-1.4B: 2.8GB download + ~11GB RAM, too large for CI", + ), +] + +# Sandwich-norm model with 96 effective layer applications at fp32; allow a +# wider op-order noise floor on GH Actions macOS-arm64 (same idiom as the +# SmolLM3 and general bridge-vs-HF parity tests). +_MACOS_ARM64 = platform.system() == "Darwin" and platform.machine() == "arm64" +FP32_NOISE_TOL = 1e-2 if _MACOS_ARM64 else 1e-3 + + +@pytest.fixture(scope="module") +def bridge() -> TransformerBridge: + return TransformerBridge.boot_transformers( + MODEL, device="cpu", dtype=torch.float32, trust_remote_code=True + ) + + +@pytest.fixture(scope="module") +def hf_eager(bridge: TransformerBridge) -> torch.nn.Module: + """HF model loaded independently of the bridge's wrapped instance. + + Depends on the bridge fixture only for ordering: booting the bridge runs + the adapter's prepare_loading patch (restoring the "default" RoPE init + that transformers v5 removed) on the cached modeling_ouro module, and + this plain from_pretrained load needs that patch too. The weights are + still a separate copy. + + The pad_token_id shim mirrors the bridge load path + (sources/transformers.py): transformers v5 no longer materializes a + default pad_token_id on configs, but Ouro's remote code reads + config.pad_token_id unconditionally. Mirroring the same eos fallback + keeps both copies' embedding padding_idx identical. + """ + config = AutoConfig.from_pretrained(MODEL, trust_remote_code=True) + if "pad_token_id" not in config.__dict__: + fallback_pad = getattr(config, "eos_token_id", None) + if isinstance(fallback_pad, list): + fallback_pad = fallback_pad[0] if fallback_pad else None + config.pad_token_id = fallback_pad + return AutoModelForCausalLM.from_pretrained( + MODEL, + config=config, + torch_dtype=torch.float32, + attn_implementation="eager", + trust_remote_code=True, + ).eval() + + +@pytest.fixture(scope="module") +def tokens(hf_eager: torch.nn.Module) -> torch.Tensor: + """A short, deterministic random token sequence. + + Random ids keep the test tokenizer-free and cover the vocab beyond the + few ids a short prompt would produce. + """ + generator = torch.Generator().manual_seed(0) + return torch.randint(0, hf_eager.config.vocab_size, (1, 16), generator=generator) + + +class TestOuroBridgeCreation: + """The bridge loads the remote-code model and wires the sandwich norms.""" + + def test_boot_transformers_succeeds(self, bridge: TransformerBridge) -> None: + assert bridge is not None + + def test_block_count_is_physical_layers( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module + ) -> None: + """n_layers counts the 24 physical layers, not the 96 loop applications.""" + assert len(bridge.blocks) == hf_eager.config.num_hidden_layers + assert bridge.cfg.n_layers == hf_eager.config.num_hidden_layers + + def test_config_flags(self, bridge: TransformerBridge) -> None: + assert bridge.cfg.normalization_type == "RMS" + assert bridge.cfg.positional_embedding_type == "rotary" + assert bridge.cfg.gated_mlp is True + + def test_sandwich_norms_wired(self, bridge: TransformerBridge) -> None: + """All four per-layer RMSNorms are bridged.""" + block = bridge.blocks[0] + for ln_name in ("ln1", "ln1_post", "ln2", "ln2_post"): + assert hasattr(block, ln_name), f"missing {ln_name} on block 0" + + +class TestOuroHFDelegation: + """The bridge wraps the live remote-code modules in place (no copies). + + Wrapping is in-place: the HF tree slot holds the bridge component itself, + and the bridge's submodule IS the module the wrapped OuroDecoderLayer / + OuroAttention executes. Assert those identities directly. + """ + + def test_blocks_are_the_hf_tree_layers(self, bridge: TransformerBridge) -> None: + assert bridge.blocks[0] is bridge.original_model.model.layers[0] + assert type(bridge.blocks[0].original_component).__name__ == "OuroDecoderLayer" + + def test_attention_projections_wrap_live_hf_modules(self, bridge: TransformerBridge) -> None: + attn = bridge.blocks[0].attn + assert type(attn.original_component).__name__ == "OuroAttention" + assert attn.q is attn.original_component.q_proj + assert attn.k is attn.original_component.k_proj + assert attn.v is attn.original_component.v_proj + assert attn.o is attn.original_component.o_proj + + def test_sandwich_norms_wrap_live_hf_modules(self, bridge: TransformerBridge) -> None: + block = bridge.blocks[0] + layer = block.original_component + assert block.ln1 is layer.input_layernorm + assert block.ln1_post is layer.input_layernorm_2 + assert block.ln2 is layer.post_attention_layernorm + assert block.ln2_post is layer.post_attention_layernorm_2 + + +class TestOuroForwardEquivalence: + """Bridge output reproduces the HF eager reference through the full UT loop.""" + + def test_forward_logits_match_hf_eager( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module, tokens: torch.Tensor + ) -> None: + with torch.inference_mode(): + bridge_logits = bridge(tokens) + hf_logits = hf_eager(tokens).logits + max_diff = (bridge_logits - hf_logits).abs().max().item() + assert max_diff < FP32_NOISE_TOL, ( + f"Ouro bridge vs HF eager logit drift={max_diff:.2e} exceeds the " + f"fp32-noise tolerance {FP32_NOISE_TOL:.0e}." + ) + + +class TestOuroLoopedDepth: + """Pin the Universal-Transformer loop semantics the adapter documents.""" + + def test_block_hook_fires_once_per_ut_step( + self, bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + """The same physical block executes total_ut_steps times per forward. + + This is the load-bearing behavioural difference from a standard + decoder: hooks on blocks.{i} fire once per loop step, and + run_with_cache keeps the final step's value. If HF's remote code ever + changes the loop (or the bridge stops delegating it), this fails. + """ + fired: list[bool] = [] + bridge.run_with_hooks( + tokens, + fwd_hooks=[("blocks.0.hook_out", lambda value, hook: fired.append(True))], + ) + assert len(fired) == bridge.cfg.total_ut_steps + + def test_hook_out_shape_matches_residual_stream( + self, bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + captured: list[torch.Tensor] = [] + + def _capture(value, hook): + captured.append(value[0] if isinstance(value, tuple) else value) + return value + + bridge.run_with_hooks(tokens, fwd_hooks=[("blocks.0.hook_out", _capture)]) + assert captured + assert captured[-1].shape == (1, tokens.shape[1], bridge.cfg.d_model) + + def test_bridge_runs_its_own_attention_reconstruction( + self, bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + """Anti-tautology guard: prove the bridge's custom attention path executes. + + If a future refactor made the bridge delegate to HF attention directly, + the parity test above would pass trivially. Assert a bridge-specific + hook fires during the forward pass instead. + """ + fired: list[bool] = [] + bridge.run_with_hooks( + tokens, + fwd_hooks=[ + ("blocks.0.attn.hook_attn_scores", lambda value, hook: fired.append(True)), + ], + ) + assert fired, ( + "blocks.0.attn.hook_attn_scores did not fire, so the bridge no longer " + "runs its own attention reconstruction and the parity test is tautological." + ) + + +class TestOuroConfigPropagation: + """Loop-related HF config attrs surface on bridge.cfg via _HF_PASSTHROUGH_ATTRS.""" + + def test_total_ut_steps_propagates( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module + ) -> None: + assert bridge.cfg.total_ut_steps == hf_eager.config.total_ut_steps + + def test_early_exit_threshold_propagates( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module + ) -> None: + assert bridge.cfg.early_exit_threshold == hf_eager.config.early_exit_threshold diff --git a/tests/unit/model_bridge/supported_architectures/test_ouro_adapter.py b/tests/unit/model_bridge/supported_architectures/test_ouro_adapter.py new file mode 100644 index 000000000..2031ca91e --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_ouro_adapter.py @@ -0,0 +1,274 @@ +"""Unit tests for OuroArchitectureAdapter. + +Tests cover: +- Config attributes set by the adapter +- Component mapping structure and HF module paths, including Ouro's + sandwich normalization (four RMSNorms per decoder layer) +- Standard Q/K/V/O weight conversion rules +- setup_component_testing rotary embedding wiring +- Factory registration + +Ouro's Universal-Transformer loop and early-exit gate live in the remote-code +HF forward and are deliberately NOT mapped by the adapter; the top-level-keys +test pins that scope (no "gate" / per-step components). +""" + +from types import SimpleNamespace + +import pytest + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.ouro import ( + OuroArchitectureAdapter, +) + + +def _make_cfg( + n_heads: int = 4, + d_model: int = 64, + n_layers: int = 2, + d_mlp: int = 256, + d_vocab: int = 100, + n_ctx: int = 64, +) -> TransformerBridgeConfig: + # Keep dimensions tiny so adapter tests do not need HF downloads or real checkpoints. + # Ouro uses full MHA (num_key_value_heads == num_attention_heads). + return TransformerBridgeConfig( + d_model=d_model, + d_head=d_model // n_heads, + n_layers=n_layers, + n_ctx=n_ctx, + n_heads=n_heads, + n_key_value_heads=n_heads, + d_vocab=d_vocab, + d_mlp=d_mlp, + architecture="OuroForCausalLM", + ) + + +@pytest.fixture +def cfg() -> TransformerBridgeConfig: + return _make_cfg() + + +@pytest.fixture +def adapter(cfg: TransformerBridgeConfig) -> OuroArchitectureAdapter: + return OuroArchitectureAdapter(cfg) + + +def _fake_hf_model(rotary_emb: object) -> SimpleNamespace: + return SimpleNamespace(model=SimpleNamespace(rotary_emb=rotary_emb)) + + +class DummyAttention: + def __init__(self) -> None: + self.rotary_emb = None + + def set_rotary_emb(self, rotary_emb: object) -> None: + self.rotary_emb = rotary_emb + + +class DummyBlock: + def __init__(self, has_attention: bool = True) -> None: + if has_attention: + self.attn = DummyAttention() + + +class DummyBridgeModel: + def __init__(self, blocks: list[DummyBlock]) -> None: + self.blocks = blocks + + +class TestOuroAdapterConfig: + """Adapter-owned config defaults that downstream bridge code relies on.""" + + def test_normalization_flags(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "RMS" + assert adapter.cfg.uses_rms_norm is True + assert adapter.cfg.final_rms is True + + def test_rotary_and_mlp_flags(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.cfg.positional_embedding_type == "rotary" + assert adapter.cfg.gated_mlp is True + assert adapter.cfg.attn_only is False + + def test_no_rmsnorm_offset(self, adapter: OuroArchitectureAdapter) -> None: + """Ouro's RMSNorm applies the weight directly (no Gemma-style +1.0 offset).""" + assert not getattr(adapter.cfg, "rmsnorm_uses_offset", False) + + def test_supports_fold_ln_is_false(self, adapter: OuroArchitectureAdapter) -> None: + """ln_final runs after every UT pass, feeding the next pass and the exit + gate; folding it into W_U would corrupt passes 1..N-1 in the live module.""" + assert adapter.supports_fold_ln is False + + +class TestOuroComponentMapping: + """The adapter contract: TL canonical names mapped to Ouro HF module paths.""" + + def test_top_level_keys(self, adapter: OuroArchitectureAdapter) -> None: + # No "gate" key: model.early_exit_gate is intentionally unmapped. + assert set(adapter.component_mapping.keys()) == { + "embed", + "rotary_emb", + "blocks", + "ln_final", + "unembed", + } + + def test_embed_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["embed"].name == "model.embed_tokens" + + def test_rotary_emb_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["rotary_emb"].name == "model.rotary_emb" + + def test_blocks_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["blocks"].name == "model.layers" + + def test_ln_final_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["ln_final"].name == "model.norm" + + def test_unembed_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["unembed"].name == "lm_head" + + def test_block_submodule_keys(self, adapter: OuroArchitectureAdapter) -> None: + """Sandwich norm: four RMSNorms per block, not the usual two.""" + blocks = adapter.component_mapping["blocks"] + assert set(blocks.submodules.keys()) == { + "ln1", + "ln1_post", + "ln2", + "ln2_post", + "attn", + "mlp", + } + + def test_sandwich_norm_hf_paths(self, adapter: OuroArchitectureAdapter) -> None: + """The extra norms map to Ouro's *_2 module names. + + Forward order in OuroDecoderLayer: input_layernorm (pre-attn) -> attn + -> input_layernorm_2 (post-attn, pre-residual); post_attention_layernorm + (pre-MLP) -> mlp -> post_attention_layernorm_2 (post-MLP, pre-residual). + """ + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["ln1"].name == "input_layernorm" + assert blocks.submodules["ln1_post"].name == "input_layernorm_2" + assert blocks.submodules["ln2"].name == "post_attention_layernorm" + assert blocks.submodules["ln2_post"].name == "post_attention_layernorm_2" + + def test_attention_submodule_keys(self, adapter: OuroArchitectureAdapter) -> None: + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert set(attn.submodules.keys()) == {"q", "k", "v", "o"} + + def test_mlp_submodule_keys(self, adapter: OuroArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert set(mlp.submodules.keys()) == {"gate", "in", "out"} + + def test_bridge_types(self, adapter: OuroArchitectureAdapter) -> None: + mapping = adapter.component_mapping + blocks = mapping["blocks"] + assert isinstance(mapping["embed"], EmbeddingBridge) + assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge) + assert isinstance(blocks, BlockBridge) + assert isinstance(mapping["ln_final"], RMSNormalizationBridge) + assert isinstance(mapping["unembed"], UnembeddingBridge) + for ln_key in ("ln1", "ln1_post", "ln2", "ln2_post"): + assert isinstance(blocks.submodules[ln_key], RMSNormalizationBridge) + assert isinstance(blocks.submodules["attn"], PositionEmbeddingsAttentionBridge) + assert isinstance(blocks.submodules["mlp"], GatedMLPBridge) + + def test_attention_hf_paths(self, adapter: OuroArchitectureAdapter) -> None: + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert attn.name == "self_attn" + assert attn.submodules["q"].name == "q_proj" + assert attn.submodules["k"].name == "k_proj" + assert attn.submodules["v"].name == "v_proj" + assert attn.submodules["o"].name == "o_proj" + + def test_mlp_hf_paths(self, adapter: OuroArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert mlp.name == "mlp" + assert mlp.submodules["gate"].name == "gate_proj" + assert mlp.submodules["in"].name == "up_proj" + assert mlp.submodules["out"].name == "down_proj" + + def test_linear_submodule_bridge_types(self, adapter: OuroArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + attn = blocks.submodules["attn"] + mlp = blocks.submodules["mlp"] + for submodule in [*attn.submodules.values(), *mlp.submodules.values()]: + assert isinstance(submodule, LinearBridge) + + +class TestOuroWeightConversions: + """Standard split-QKV conversion rules from _qkvo_weight_conversions().""" + + def test_qkvo_conversion_keys_present(self, adapter: OuroArchitectureAdapter) -> None: + for key in ( + "blocks.{i}.attn.q.weight", + "blocks.{i}.attn.k.weight", + "blocks.{i}.attn.v.weight", + "blocks.{i}.attn.o.weight", + ): + assert key in adapter.weight_processing_conversions + + +class TestOuroSetupComponentTesting: + """setup_component_testing must wire Ouro's shared rotary embedding into attention bridges.""" + + def test_sets_rotary_emb_on_template_attention(self, adapter: OuroArchitectureAdapter) -> None: + rotary_emb = object() + attn_template = adapter.get_generalized_component("blocks.0.attn") + assert isinstance(attn_template, PositionEmbeddingsAttentionBridge) + assert attn_template._rotary_emb is None + + adapter.setup_component_testing(_fake_hf_model(rotary_emb)) + + assert attn_template._rotary_emb is rotary_emb + + def test_sets_rotary_emb_on_each_bridge_model_attention( + self, adapter: OuroArchitectureAdapter + ) -> None: + rotary_emb = object() + bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(), DummyBlock()]) + + adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model) + + for block in bridge_model.blocks: + assert block.attn.rotary_emb is rotary_emb + + def test_skips_bridge_blocks_without_attention(self, adapter: OuroArchitectureAdapter) -> None: + rotary_emb = object() + bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(has_attention=False)]) + + adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model) + + assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb + + +class TestOuroFactoryRegistration: + """The factory resolves OuroForCausalLM to this adapter.""" + + def test_supported_architectures_entry(self) -> None: + from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, + ) + + assert SUPPORTED_ARCHITECTURES["OuroForCausalLM"] is OuroArchitectureAdapter + + def test_factory_selects_ouro_adapter(self, cfg: TransformerBridgeConfig) -> None: + from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, + ) + + adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg) + assert isinstance(adapter, OuroArchitectureAdapter) diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 748c160ca..cef0d22d7 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -66,6 +66,7 @@ OlmoeArchitectureAdapter, OpenElmArchitectureAdapter, OptArchitectureAdapter, + OuroArchitectureAdapter, Phi3ArchitectureAdapter, PhiArchitectureAdapter, PhiMoEArchitectureAdapter, @@ -148,6 +149,7 @@ "OlmoeForCausalLM": OlmoeArchitectureAdapter, "OpenELMForCausalLM": OpenElmArchitectureAdapter, "OPTForCausalLM": OptArchitectureAdapter, + "OuroForCausalLM": OuroArchitectureAdapter, "PhiForCausalLM": PhiArchitectureAdapter, "Phi3ForCausalLM": Phi3ArchitectureAdapter, "PhiMoEForCausalLM": PhiMoEArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 4729f92cf..c64930a16 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -81,6 +81,9 @@ "num_mem_blocks", "layers_block_type", "use_shared_attention_adapter", + # Ouro (LoopLM) + "total_ut_steps", + "early_exit_threshold", ] diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 0ea8220af..2d67a01b9 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -265,6 +265,7 @@ def determine_architecture_from_hf_config(hf_config): "qwen3_5_text": "Qwen3_5ForCausalLM", "smollm3": "SmolLM3ForCausalLM", "openelm": "OpenELMForCausalLM", + "ouro": "OuroForCausalLM", "stablelm": "StableLmForCausalLM", "t5": "T5ForConditionalGeneration", "mt5": "MT5ForConditionalGeneration", @@ -581,6 +582,9 @@ def boot( "num_mem_blocks", "layers_block_type", "use_shared_attention_adapter", + # Ouro (LoopLM) + "total_ut_steps", + "early_exit_threshold", ] for attr in _HF_PASSTHROUGH_ATTRS: val = getattr(hf_config, attr, None) diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 0f0e2fd0b..6ab2ad971 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -169,6 +169,9 @@ from transformer_lens.model_bridge.supported_architectures.opt import ( OptArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.ouro import ( + OuroArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.phi import ( PhiArchitectureAdapter, ) @@ -282,6 +285,7 @@ "Olmo3ArchitectureAdapter", "OlmoeArchitectureAdapter", "OptArchitectureAdapter", + "OuroArchitectureAdapter", "PhiArchitectureAdapter", "Phi3ArchitectureAdapter", "PhiMoEArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/ouro.py b/transformer_lens/model_bridge/supported_architectures/ouro.py new file mode 100644 index 000000000..e443c49f9 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/ouro.py @@ -0,0 +1,227 @@ +"""Ouro architecture adapter.""" + +import sys +from typing import Any, Optional + +import torch + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) + + +def _compute_default_rope_parameters( + config: Any, + device: Optional[torch.device] = None, + seq_len: Optional[int] = None, + **rope_kwargs: Any, +) -> tuple[torch.Tensor, float]: + """Standard (unscaled) RoPE inverse frequencies, as transformers v4 defined them. + + Transformers v5 removed the "default" entry from ROPE_INIT_FUNCTIONS (standard + RoPE moved to a per-model static method), but Ouro's remote code still looks it + up. Signature and return match the v4 contract the remote code calls with: + (config, device) -> (inv_freq, attention_scaling). + """ + base = config.rope_theta + partial_rotary_factor = getattr(config, "partial_rotary_factor", None) or 1.0 + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / ( + base + ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, 1.0 + + +class OuroArchitectureAdapter(ArchitectureAdapter): + """Architecture adapter for ByteDance Ouro (LoopLM) models. + + Ouro is a looped-depth ("Universal Transformer") decoder: the remote-code + ``OuroModel.forward`` applies the same ``num_hidden_layers``-deep stack + ``total_ut_steps`` times (4 for the released checkpoints) within a single + forward pass, applying ``model.norm`` after every pass. The loop lives + entirely inside the HF forward, which the bridge delegates to, so logits + and generation are correct with no loop handling here. ``n_layers`` counts + the physical layers; each block's hooks fire once per loop step, and a + cache records the final step's value. The same holds for ``ln_final`` + (``model.norm``): it runs after EVERY UT pass, so its hooks fire + ``total_ut_steps`` times per forward and ``run_with_cache`` keeps only the + last pass. + + The backbone is Qwen2/Llama-shaped (RoPE, no-bias q/k/v/o projections, + SwiGLU gate/up/down MLP, untied lm_head) with one twist: sandwich + normalization. Each decoder layer has FOUR RMSNorms; the extra two + (``input_layernorm_2``, ``post_attention_layernorm_2``) apply to the + sublayer outputs before the residual add, exactly like Gemma2's + ``ln1_post``/``ln2_post`` but without Gemma's +1.0 RMSNorm offset. + + Deliberately not mapped by this adapter: + + - per-loop-step hooks (a cache holds the final UT step only) + - ``model.early_exit_gate``, the adaptive-exit halting head + - the ``UniversalTransformerCache`` slot layout (``step * n_layers + layer``) + + Loading requires ``trust_remote_code=True`` (``auto_map`` to + ``modeling_ouro``). + + Optional Parameters (may not exist in state_dict): + ------------------------------------------------- + Ouro models do NOT have biases on any mapped linear layers: + + - blocks.{i}.attn.b_Q / b_K / b_V / b_O - no attention biases + - blocks.{i}.mlp.b_gate / b_in / b_out - no MLP biases + - blocks.{i}.ln1.b / ln1_post.b / ln2.b / ln2_post.b - RMSNorm has no bias + - ln_final.b - RMSNorm has no bias + + Weight processing must handle these missing biases gracefully using + ProcessWeights._safe_get_tensor() or by checking for None values. + """ + + def __init__(self, cfg: Any) -> None: + """Initialize the Ouro architecture adapter.""" + super().__init__(cfg) + + # Set config variables for weight processing + 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.cfg.uses_rms_norm = True + # default_prepend_bos stays at the framework default: the GPT2-style BPE + # tokenizer (bos == eos == <|endoftext|>) does not prepend BOS itself. + + # ln_final (model.norm) is applied after EVERY UT pass, feeding the next + # pass and the early-exit gate, so it is not a final-only norm. Folding + # it into W_U resets the live module's norm weight the loop reuses and + # corrupts UT passes 1..N-1. + self.supports_fold_ln = False + + self.weight_processing_conversions = { + **self._qkvo_weight_conversions(), + } + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), + "blocks": BlockBridge( + name="model.layers", + config=self.cfg, + submodules={ + "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), + "ln1_post": RMSNormalizationBridge(name="input_layernorm_2", config=self.cfg), + "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), + "ln2_post": RMSNormalizationBridge( + name="post_attention_layernorm_2", config=self.cfg + ), + "attn": PositionEmbeddingsAttentionBridge( + name="self_attn", + config=self.cfg, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="o_proj"), + }, + requires_attention_mask=True, + requires_position_embeddings=True, + ), + "mlp": GatedMLPBridge( + name="mlp", + config=self.cfg, + submodules={ + "gate": LinearBridge(name="gate_proj"), + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), + } + + def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: + """Patch Ouro's remote code for compatibility with transformers v5. + + Ouro's modeling code was written against transformers 4.55, where + standard RoPE lived in ROPE_INIT_FUNCTIONS["default"]. Transformers v5 + removed that key and instead expects each *RotaryEmbedding class to + carry a compute_default_rope_parameters static method. Two call sites + break, so two patches: + + 1. OuroRotaryEmbedding.__init__ does ROPE_INIT_FUNCTIONS["default"] + (KeyError). Rebind the module-level name inside the imported + modeling_ouro module(s) to a copy with "default" restored; the + shared transformers dict is left untouched. + 2. v5's PreTrainedModel._init_weights re-initializes RotaryEmbedding + buffers via module.compute_default_rope_parameters(config) + (AttributeError). Attach the same function as a static method. + + Args: + model_name: The HuggingFace model name/path + model_kwargs: The kwargs dict for from_pretrained() + """ + # Force-import the modeling module so we can patch it + try: + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + get_class_from_dynamic_module( + "modeling_ouro.OuroForCausalLM", + model_name, + ) + except Exception: + return + + # Each checkpoint revision gets its own module in sys.modules, so patch + # every imported Ouro modeling module (same idiom as openelm.py). + for key in list(sys.modules.keys()): + if "ouro" in key.lower() and "modeling" in key.lower(): + module = sys.modules[key] + rope_functions = getattr(module, "ROPE_INIT_FUNCTIONS", None) + if rope_functions is not None and "default" not in rope_functions: + setattr( + module, + "ROPE_INIT_FUNCTIONS", + {**rope_functions, "default": _compute_default_rope_parameters}, + ) + rope_class = getattr(module, "OuroRotaryEmbedding", None) + if rope_class is not None and not hasattr( + rope_class, "compute_default_rope_parameters" + ): + rope_class.compute_default_rope_parameters = staticmethod( + _compute_default_rope_parameters + ) + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Ouro component testing. + + Ouro uses RoPE (Rotary Position Embeddings) with a single shared + ``model.rotary_emb``. We set the rotary_emb reference on all attention + bridge instances for component testing. + + Args: + hf_model: The HuggingFace Ouro model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from the model + rotary_emb = hf_model.model.rotary_emb + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 9a38c22b5..a304771ef 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -99,6 +99,7 @@ "OlmoForCausalLM", "OlmoeForCausalLM", "OPTForCausalLM", + "OuroForCausalLM", "PhiForCausalLM", "Phi3ForCausalLM", "PhiMoEForCausalLM", @@ -178,6 +179,7 @@ "OlmoForCausalLM": ["allenai"], "OpenELMForCausalLM": ["apple"], "OPTForCausalLM": ["facebook"], + "OuroForCausalLM": ["ByteDance"], "Phi3ForCausalLM": ["microsoft"], "PhiMoEForCausalLM": ["microsoft"], "PhiForCausalLM": ["microsoft"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index 399603b39..b42d76ff9 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -6,9 +6,9 @@ "min_downloads": 500, "scan_duration_seconds": 8.1 }, - "total_architectures": 69, - "total_models": 13144, - "total_verified": 1028, + "total_architectures": 70, + "total_models": 13145, + "total_verified": 1029, "models": [ { "architecture_id": "FalconH1ForCausalLM", @@ -156373,6 +156373,20 @@ "phase7_score": null, "phase8_score": null }, + { + "architecture_id": "OuroForCausalLM", + "model_id": "ByteDance/Ouro-1.4B", + "status": 1, + "verified_date": "2026-07-07", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 95.9, + "phase7_score": null, + "phase8_score": null + }, { "architecture_id": "SmolLM3ForCausalLM", "model_id": "HuggingFaceTB/SmolLM3-3B", diff --git a/transformer_lens/tools/model_registry/data/verification_history.json b/transformer_lens/tools/model_registry/data/verification_history.json index 530f03e15..54eb14cfb 100644 --- a/transformer_lens/tools/model_registry/data/verification_history.json +++ b/transformer_lens/tools/model_registry/data/verification_history.json @@ -1,5 +1,5 @@ { - "last_updated": "2026-07-02T02:19:51.172022", + "last_updated": "2026-07-07T23:04:00.376842", "records": [ { "model_id": "Macropodus/macbert4mdcspell_v1", @@ -16860,6 +16860,16 @@ "notes": "Full verification completed", "invalidated": false, "invalidation_reason": null + }, + { + "model_id": "ByteDance/Ouro-1.4B", + "architecture_id": "OuroForCausalLM", + "verified_date": "2026-07-07", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null } ] } diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index d62edda90..6f514dd6c 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -66,6 +66,7 @@ "BD3LM": "Kuleshov Group's Block Diffusion Language Model (ICLR 2025) for masked text generation", "HunYuanDenseV1ForCausalLM": "Tencent's open source decoder models", "Cohere2ForCausalLM": "Cohere's Command-A architecture with interleaved sliding-window RoPE and full-attention NoPE layers", + "OuroForCausalLM": "ByteDance's Ouro looped language model (LoopLM) with weight-shared iterated depth", # Unsupported architectures "BertModel": "Google's BERT bidirectional encoder for understanding tasks", "BertForMaskedLM": "BERT with masked language modeling head", diff --git a/transformer_lens/tools/model_registry/verify_models.py b/transformer_lens/tools/model_registry/verify_models.py index 407acc975..667daf4e6 100644 --- a/transformer_lens/tools/model_registry/verify_models.py +++ b/transformer_lens/tools/model_registry/verify_models.py @@ -62,6 +62,7 @@ # These are not in the legacy NEED_REMOTE_CODE_MODELS tuple (loading_from_pretrained.py). _BRIDGE_REMOTE_CODE_PREFIXES: tuple[str, ...] = ( "baichuan-inc/", # BaichuanForCausalLM — ships own modeling_baichuan.py + "ByteDance/Ouro-", # OuroForCausalLM — ships own modeling_ouro.py "internlm/", # InternLM2ForCausalLM — ships own modeling_internlm2.py "kuleshov-group/", # BD3LM — ships own custom modeling_d_dit.py )