transformer_lens.model_bridge.supported_architectures.ouro module

Ouro architecture adapter.

class transformer_lens.model_bridge.supported_architectures.ouro.OuroArchitectureAdapter(cfg: Any)

Bases: 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.

__init__(cfg: Any) → None

Initialize the Ouro architecture adapter.

prepare_loading(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.

Parameters:

  • model_name – The HuggingFace model name/path

  • model_kwargs – The kwargs dict for from_pretrained()

setup_component_testing(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.

Parameters:

  • hf_model – The HuggingFace Ouro model instance

  • bridge_model – The TransformerBridge model (if available, set rotary_emb on actual instances)