Coverage for transformer_lens/model_bridge/supported_architectures/ouro.py: 43%
22 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Ouro architecture adapter."""
3from typing import Any
5from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
6from transformer_lens.model_bridge.generalized_components import (
7 BlockBridge,
8 EmbeddingBridge,
9 LinearBridge,
10 PositionEmbeddingsAttentionBridge,
11 RMSNormalizationBridge,
12 RotaryEmbeddingBridge,
13 UnembeddingBridge,
14)
15from transformer_lens.model_bridge.supported_architectures._remote_code_compat import (
16 compute_default_rope_inv_freq,
17 force_import_remote_class,
18 iter_remote_modeling_modules,
19)
22class OuroArchitectureAdapter(ArchitectureAdapter):
23 """Architecture adapter for ByteDance Ouro (LoopLM) models.
25 Ouro is a looped-depth ("Universal Transformer") decoder: the remote-code
26 ``OuroModel.forward`` applies the same ``num_hidden_layers``-deep stack
27 ``total_ut_steps`` times (4 for the released checkpoints) within a single
28 forward pass, applying ``model.norm`` after every pass. The loop lives
29 entirely inside the HF forward, which the bridge delegates to, so logits
30 and generation are correct with no loop handling here. ``n_layers`` counts
31 the physical layers; each block's hooks fire once per loop step, and a
32 cache records the final step's value. The same holds for ``ln_final``
33 (``model.norm``): it runs after EVERY UT pass, so its hooks fire
34 ``total_ut_steps`` times per forward and ``run_with_cache`` keeps only the
35 last pass.
37 The backbone is Qwen2/Llama-shaped (RoPE, no-bias q/k/v/o projections,
38 SwiGLU gate/up/down MLP, untied lm_head) with one twist: sandwich
39 normalization. Each decoder layer has FOUR RMSNorms; the extra two
40 (``input_layernorm_2``, ``post_attention_layernorm_2``) apply to the
41 sublayer outputs before the residual add, exactly like Gemma2's
42 ``ln1_post``/``ln2_post`` but without Gemma's +1.0 RMSNorm offset.
44 Deliberately not mapped by this adapter:
46 - per-loop-step hooks (a cache holds the final UT step only)
47 - ``model.early_exit_gate``, the adaptive-exit halting head
48 - the ``UniversalTransformerCache`` slot layout (``step * n_layers + layer``)
50 Loading requires ``trust_remote_code=True`` (``auto_map`` to
51 ``modeling_ouro``).
53 Optional Parameters (may not exist in state_dict):
54 -------------------------------------------------
55 Ouro models do NOT have biases on any mapped linear layers:
57 - blocks.{i}.attn.b_Q / b_K / b_V / b_O - no attention biases
58 - blocks.{i}.mlp.b_gate / b_in / b_out - no MLP biases
59 - blocks.{i}.ln1.b / ln1_post.b / ln2.b / ln2_post.b - RMSNorm has no bias
60 - ln_final.b - RMSNorm has no bias
62 Weight processing must handle these missing biases gracefully using
63 ProcessWeights._safe_get_tensor() or by checking for None values.
64 """
66 _testing_eager = None
68 def __init__(self, cfg: Any) -> None:
69 """Initialize the Ouro architecture adapter."""
70 super().__init__(cfg)
72 self._set_rms_rotary_defaults()
73 # default_prepend_bos stays at the framework default: the GPT2-style BPE
74 # tokenizer (bos == eos == <|endoftext|>) does not prepend BOS itself.
76 # ln_final (model.norm) is applied after EVERY UT pass, feeding the next
77 # pass and the early-exit gate, so it is not a final-only norm. Folding
78 # it into W_U resets the live module's norm weight the loop reuses and
79 # corrupts UT passes 1..N-1.
80 self.supports_fold_ln = False
82 self.weight_processing_conversions = {
83 **self._qkvo_weight_conversions(),
84 }
85 self.component_mapping = {
86 "embed": EmbeddingBridge(name="model.embed_tokens"),
87 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"),
88 "blocks": BlockBridge(
89 name="model.layers",
90 config=self.cfg,
91 submodules={
92 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
93 "ln1_post": RMSNormalizationBridge(name="input_layernorm_2", config=self.cfg),
94 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
95 "ln2_post": RMSNormalizationBridge(
96 name="post_attention_layernorm_2", config=self.cfg
97 ),
98 "attn": PositionEmbeddingsAttentionBridge(
99 name="self_attn",
100 config=self.cfg,
101 submodules={
102 "q": LinearBridge(name="q_proj"),
103 "k": LinearBridge(name="k_proj"),
104 "v": LinearBridge(name="v_proj"),
105 "o": LinearBridge(name="o_proj"),
106 },
107 requires_attention_mask=True,
108 requires_position_embeddings=True,
109 ),
110 "mlp": self._gated_mlp(),
111 },
112 ),
113 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
114 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
115 }
117 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
118 """Patch Ouro's remote code for compatibility with transformers v5.
120 Ouro's modeling code was written against transformers 4.55, where
121 standard RoPE lived in ROPE_INIT_FUNCTIONS["default"]. Transformers v5
122 removed that key and instead expects each *RotaryEmbedding class to
123 carry a compute_default_rope_parameters static method. Two call sites
124 break, so two patches:
126 1. OuroRotaryEmbedding.__init__ does ROPE_INIT_FUNCTIONS["default"]
127 (KeyError). Rebind the module-level name inside the imported
128 modeling_ouro module(s) to a copy with "default" restored; the
129 shared transformers dict is left untouched.
130 2. v5's PreTrainedModel._init_weights re-initializes RotaryEmbedding
131 buffers via module.compute_default_rope_parameters(config)
132 (AttributeError). Attach the same function as a static method.
134 Args:
135 model_name: The HuggingFace model name/path
136 model_kwargs: The kwargs dict for from_pretrained()
137 """
138 # Force-import the modeling module so we can patch it
139 if force_import_remote_class(model_name, "modeling_ouro.OuroForCausalLM") is None:
140 return
142 # Each checkpoint revision gets its own module in sys.modules; patch all.
143 for module in iter_remote_modeling_modules("ouro"):
144 # Rebind the module-level ROPE_INIT_FUNCTIONS name to a copy with
145 # "default" restored; the shared transformers dict stays untouched.
146 rope_functions = getattr(module, "ROPE_INIT_FUNCTIONS", None)
147 if rope_functions is not None and "default" not in rope_functions:
148 setattr(
149 module,
150 "ROPE_INIT_FUNCTIONS",
151 {**rope_functions, "default": compute_default_rope_inv_freq},
152 )
153 rope_class = getattr(module, "OuroRotaryEmbedding", None)
154 if rope_class is not None and not hasattr(
155 rope_class, "compute_default_rope_parameters"
156 ):
157 rope_class.compute_default_rope_parameters = staticmethod(
158 compute_default_rope_inv_freq
159 )