Coverage for transformer_lens/model_bridge/supported_architectures/ouro.py: 57%
47 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""Ouro architecture adapter."""
3import sys
4from typing import Any, Optional
6import torch
8from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
9from transformer_lens.model_bridge.generalized_components import (
10 BlockBridge,
11 EmbeddingBridge,
12 GatedMLPBridge,
13 LinearBridge,
14 PositionEmbeddingsAttentionBridge,
15 RMSNormalizationBridge,
16 RotaryEmbeddingBridge,
17 UnembeddingBridge,
18)
21def _compute_default_rope_parameters(
22 config: Any,
23 device: Optional[torch.device] = None,
24 seq_len: Optional[int] = None,
25 **rope_kwargs: Any,
26) -> tuple[torch.Tensor, float]:
27 """Standard (unscaled) RoPE inverse frequencies, as transformers v4 defined them.
29 Transformers v5 removed the "default" entry from ROPE_INIT_FUNCTIONS (standard
30 RoPE moved to a per-model static method), but Ouro's remote code still looks it
31 up. Signature and return match the v4 contract the remote code calls with:
32 (config, device) -> (inv_freq, attention_scaling).
33 """
34 base = config.rope_theta
35 partial_rotary_factor = getattr(config, "partial_rotary_factor", None) or 1.0
36 head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
37 dim = int(head_dim * partial_rotary_factor)
38 inv_freq = 1.0 / (
39 base
40 ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
41 )
42 return inv_freq, 1.0
45class OuroArchitectureAdapter(ArchitectureAdapter):
46 """Architecture adapter for ByteDance Ouro (LoopLM) models.
48 Ouro is a looped-depth ("Universal Transformer") decoder: the remote-code
49 ``OuroModel.forward`` applies the same ``num_hidden_layers``-deep stack
50 ``total_ut_steps`` times (4 for the released checkpoints) within a single
51 forward pass, applying ``model.norm`` after every pass. The loop lives
52 entirely inside the HF forward, which the bridge delegates to, so logits
53 and generation are correct with no loop handling here. ``n_layers`` counts
54 the physical layers; each block's hooks fire once per loop step, and a
55 cache records the final step's value. The same holds for ``ln_final``
56 (``model.norm``): it runs after EVERY UT pass, so its hooks fire
57 ``total_ut_steps`` times per forward and ``run_with_cache`` keeps only the
58 last pass.
60 The backbone is Qwen2/Llama-shaped (RoPE, no-bias q/k/v/o projections,
61 SwiGLU gate/up/down MLP, untied lm_head) with one twist: sandwich
62 normalization. Each decoder layer has FOUR RMSNorms; the extra two
63 (``input_layernorm_2``, ``post_attention_layernorm_2``) apply to the
64 sublayer outputs before the residual add, exactly like Gemma2's
65 ``ln1_post``/``ln2_post`` but without Gemma's +1.0 RMSNorm offset.
67 Deliberately not mapped by this adapter:
69 - per-loop-step hooks (a cache holds the final UT step only)
70 - ``model.early_exit_gate``, the adaptive-exit halting head
71 - the ``UniversalTransformerCache`` slot layout (``step * n_layers + layer``)
73 Loading requires ``trust_remote_code=True`` (``auto_map`` to
74 ``modeling_ouro``).
76 Optional Parameters (may not exist in state_dict):
77 -------------------------------------------------
78 Ouro models do NOT have biases on any mapped linear layers:
80 - blocks.{i}.attn.b_Q / b_K / b_V / b_O - no attention biases
81 - blocks.{i}.mlp.b_gate / b_in / b_out - no MLP biases
82 - blocks.{i}.ln1.b / ln1_post.b / ln2.b / ln2_post.b - RMSNorm has no bias
83 - ln_final.b - RMSNorm has no bias
85 Weight processing must handle these missing biases gracefully using
86 ProcessWeights._safe_get_tensor() or by checking for None values.
87 """
89 def __init__(self, cfg: Any) -> None:
90 """Initialize the Ouro architecture adapter."""
91 super().__init__(cfg)
93 # Set config variables for weight processing
94 self.cfg.normalization_type = "RMS"
95 self.cfg.positional_embedding_type = "rotary"
96 self.cfg.final_rms = True
97 self.cfg.gated_mlp = True
98 self.cfg.attn_only = False
99 self.cfg.uses_rms_norm = True
100 # default_prepend_bos stays at the framework default: the GPT2-style BPE
101 # tokenizer (bos == eos == <|endoftext|>) does not prepend BOS itself.
103 # ln_final (model.norm) is applied after EVERY UT pass, feeding the next
104 # pass and the early-exit gate, so it is not a final-only norm. Folding
105 # it into W_U resets the live module's norm weight the loop reuses and
106 # corrupts UT passes 1..N-1.
107 self.supports_fold_ln = False
109 self.weight_processing_conversions = {
110 **self._qkvo_weight_conversions(),
111 }
112 self.component_mapping = {
113 "embed": EmbeddingBridge(name="model.embed_tokens"),
114 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"),
115 "blocks": BlockBridge(
116 name="model.layers",
117 config=self.cfg,
118 submodules={
119 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
120 "ln1_post": RMSNormalizationBridge(name="input_layernorm_2", config=self.cfg),
121 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
122 "ln2_post": RMSNormalizationBridge(
123 name="post_attention_layernorm_2", config=self.cfg
124 ),
125 "attn": PositionEmbeddingsAttentionBridge(
126 name="self_attn",
127 config=self.cfg,
128 submodules={
129 "q": LinearBridge(name="q_proj"),
130 "k": LinearBridge(name="k_proj"),
131 "v": LinearBridge(name="v_proj"),
132 "o": LinearBridge(name="o_proj"),
133 },
134 requires_attention_mask=True,
135 requires_position_embeddings=True,
136 ),
137 "mlp": GatedMLPBridge(
138 name="mlp",
139 config=self.cfg,
140 submodules={
141 "gate": LinearBridge(name="gate_proj"),
142 "in": LinearBridge(name="up_proj"),
143 "out": LinearBridge(name="down_proj"),
144 },
145 ),
146 },
147 ),
148 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
149 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
150 }
152 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
153 """Patch Ouro's remote code for compatibility with transformers v5.
155 Ouro's modeling code was written against transformers 4.55, where
156 standard RoPE lived in ROPE_INIT_FUNCTIONS["default"]. Transformers v5
157 removed that key and instead expects each *RotaryEmbedding class to
158 carry a compute_default_rope_parameters static method. Two call sites
159 break, so two patches:
161 1. OuroRotaryEmbedding.__init__ does ROPE_INIT_FUNCTIONS["default"]
162 (KeyError). Rebind the module-level name inside the imported
163 modeling_ouro module(s) to a copy with "default" restored; the
164 shared transformers dict is left untouched.
165 2. v5's PreTrainedModel._init_weights re-initializes RotaryEmbedding
166 buffers via module.compute_default_rope_parameters(config)
167 (AttributeError). Attach the same function as a static method.
169 Args:
170 model_name: The HuggingFace model name/path
171 model_kwargs: The kwargs dict for from_pretrained()
172 """
173 # Force-import the modeling module so we can patch it
174 try:
175 from transformers.dynamic_module_utils import get_class_from_dynamic_module
177 get_class_from_dynamic_module(
178 "modeling_ouro.OuroForCausalLM",
179 model_name,
180 )
181 except Exception:
182 return
184 # Each checkpoint revision gets its own module in sys.modules, so patch
185 # every imported Ouro modeling module (same idiom as openelm.py).
186 for key in list(sys.modules.keys()):
187 if "ouro" in key.lower() and "modeling" in key.lower():
188 module = sys.modules[key]
189 rope_functions = getattr(module, "ROPE_INIT_FUNCTIONS", None)
190 if rope_functions is not None and "default" not in rope_functions:
191 setattr(
192 module,
193 "ROPE_INIT_FUNCTIONS",
194 {**rope_functions, "default": _compute_default_rope_parameters},
195 )
196 rope_class = getattr(module, "OuroRotaryEmbedding", None)
197 if rope_class is not None and not hasattr(
198 rope_class, "compute_default_rope_parameters"
199 ):
200 rope_class.compute_default_rope_parameters = staticmethod(
201 _compute_default_rope_parameters
202 )
204 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
205 """Set up rotary embedding references for Ouro component testing.
207 Ouro uses RoPE (Rotary Position Embeddings) with a single shared
208 ``model.rotary_emb``. We set the rotary_emb reference on all attention
209 bridge instances for component testing.
211 Args:
212 hf_model: The HuggingFace Ouro model instance
213 bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances)
214 """
215 # Get rotary embedding instance from the model
216 rotary_emb = hf_model.model.rotary_emb
218 # Set rotary_emb on actual bridge instances in bridge_model if available
219 if bridge_model is not None and hasattr(bridge_model, "blocks"):
220 # Set on each layer's actual attention bridge instance
221 for block in bridge_model.blocks:
222 if hasattr(block, "attn"):
223 block.attn.set_rotary_emb(rotary_emb)
225 # Also set on the template for get_generalized_component() calls
226 attn_bridge = self.get_generalized_component("blocks.0.attn")
227 attn_bridge.set_rotary_emb(rotary_emb)