Coverage for transformer_lens/model_bridge/supported_architectures/hrm_text.py: 48%
58 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""HRM-Text architecture adapter.
3HRM-Text (Sapient Intelligence) is a hierarchical two-timescale recurrent model:
4two transformer stacks (H = slow/planning, L = fast/computation) iterate in a
5nested loop with additive cross-stack coupling.
7Architecture notes:
8 - **Two physical stacks**: ``model.L_module`` and ``model.H_module``, each with
9 ``num_layers_per_stack`` layers. The stacks share identical internal structure
10 but have separate weights.
11 - **Recurrence**: outer H-cycle iterates ``H_cycles`` times; each iteration runs
12 ``L_cycles`` inner L-cycle iterations. Total forward passes through the layer
13 stacks = ``H_cycles * (L_cycles + 1)``. The config field ``num_hidden_layers``
14 is rewritten by HF to ``num_layers_per_stack * H_cycles * (L_cycles + 1)`` to
15 size the KV cache slots.
16 - **Parameterless RMSNorm**: ``input_layernorm``, ``post_attention_layernorm``,
17 and each stack's ``final_norm`` have no learnable weight tensor.
18 - **Sigmoid attention gate**: each attention block has a ``gate_proj`` linear
19 that produces a per-head sigmoid gate applied to the attention output before
20 ``o_proj``. Delegated to HF; hookable via ``L_blocks.{i}.attn.gate.hook_out``.
21 - **Embedding scale**: ``inputs_embeds *= embedding_scale`` (default ~39.19 for
22 HRM-Text-1B). Applied at runtime by ``HrmTextModel.forward``; must NOT be
23 folded into ``embed.weight`` — same reasoning as ``gemma1.py``.
24 - **PrefixLM mask**: instruction tokens attend bidirectionally when
25 ``token_type_ids`` is passed to HF forward; delegated, not modeled by bridge.
27Known limitations:
28 1. Hooks on ``L_blocks.{i}.*`` fire ``H_cycles * L_cycles`` times per forward;
29 on ``H_blocks.{i}.*`` they fire ``H_cycles`` times. No per-iteration index is
30 exposed; per-cycle disambiguation is future work.
31 2. Compat-mode with PrefixLM (``token_type_ids``) inputs is untested in v1.
32 3. ``supports_fold_ln = False`` — parameterless norms cannot be folded.
33 4. ``supports_center_writing_weights = False`` — block naming (``L_blocks`` /
34 ``H_blocks`` instead of ``blocks``) is incompatible with weight-centering
35 iteration over ``range(cfg.n_layers)``.
36 5. Requires ``transformers >= 5.9.0`` at runtime.
37"""
39from typing import Any
41from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
42from transformer_lens.conversion_utils.param_processing_conversion import (
43 ParamProcessingConversion,
44)
45from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
46from transformer_lens.model_bridge.generalized_components import (
47 BlockBridge,
48 EmbeddingBridge,
49 GatedMLPBridge,
50 LinearBridge,
51 PositionEmbeddingsAttentionBridge,
52 RMSNormalizationBridge,
53 RotaryEmbeddingBridge,
54 UnembeddingBridge,
55)
58class HrmTextArchitectureAdapter(ArchitectureAdapter):
59 """Architecture adapter for HRM-Text (Sapient Intelligence).
61 Exposes ``L_blocks`` (fast/low-level stack) and ``H_blocks`` (slow/high-level
62 stack) as sibling block lists. The nested recurrence loop is owned by HF's
63 forward; hooks fire once per iteration through the physical layers.
64 """
66 supports_fold_ln = False
67 supports_center_writing_weights = False
68 applicable_phases = [1, 2, 3]
70 def __init__(self, cfg: Any) -> None:
71 """Initialize the HRM-Text architecture adapter."""
72 super().__init__(cfg)
74 self.cfg.normalization_type = "RMS"
75 self.cfg.positional_embedding_type = "rotary"
76 self.cfg.final_rms = True
77 self.cfg.gated_mlp = True
78 self.cfg.attn_only = False
79 self.cfg.uses_rms_norm = True
81 if hasattr(cfg, "num_key_value_heads") and cfg.num_key_value_heads is not None: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true
82 self.cfg.n_key_value_heads = cfg.num_key_value_heads
83 elif hasattr(cfg, "num_attention_heads"): 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true
84 self.cfg.n_key_value_heads = cfg.num_attention_heads
86 for attr in (
87 "H_cycles",
88 "L_cycles",
89 "L_bp_cycles",
90 "num_layers_per_stack",
91 "embedding_scale",
92 "prefix_lm",
93 ):
94 if hasattr(cfg, attr):
95 setattr(self.cfg, attr, getattr(cfg, attr))
97 n_kv_heads = (
98 self.cfg.n_key_value_heads
99 if hasattr(self.cfg, "n_key_value_heads") and self.cfg.n_key_value_heads is not None
100 else self.cfg.n_heads
101 )
102 self.weight_processing_conversions = self._build_weight_conversions(n_kv_heads)
104 def _make_block_submodules():
105 return {
106 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
107 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
108 "attn": PositionEmbeddingsAttentionBridge(
109 name="self_attn",
110 config=self.cfg,
111 submodules={
112 "q": LinearBridge(name="q_proj"),
113 "k": LinearBridge(name="k_proj"),
114 "v": LinearBridge(name="v_proj"),
115 "o": LinearBridge(name="o_proj"),
116 "gate": LinearBridge(name="gate_proj"),
117 },
118 requires_attention_mask=True,
119 requires_position_embeddings=True,
120 ),
121 "mlp": GatedMLPBridge(
122 name="mlp",
123 config=self.cfg,
124 submodules={
125 "gate": LinearBridge(name="gate_proj"),
126 "in": LinearBridge(name="up_proj"),
127 "out": LinearBridge(name="down_proj"),
128 },
129 ),
130 }
132 self.component_mapping = {
133 "embed": EmbeddingBridge(name="model.embed_tokens"),
134 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
135 "L_blocks": BlockBridge(
136 name="model.L_module.layers",
137 submodules=_make_block_submodules(),
138 ),
139 "H_blocks": BlockBridge(
140 name="model.H_module.layers",
141 submodules=_make_block_submodules(),
142 ),
143 "L_ln_final": RMSNormalizationBridge(name="model.L_module.final_norm", config=self.cfg),
144 "H_ln_final": RMSNormalizationBridge(name="model.H_module.final_norm", config=self.cfg),
145 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
146 }
148 def _build_weight_conversions(
149 self, n_kv_heads: int
150 ) -> dict[str, ParamProcessingConversion | str]:
151 """Build weight processing conversions for both L and H block stacks.
153 Each Q/K/V/O weight under ``L_blocks.{i}`` and ``H_blocks.{i}`` needs
154 the same ``(n_heads * d_head, d_model) → (n_heads, d_head, d_model)``
155 rearrangement as a standard decoder adapter, but with the ``L_blocks`` /
156 ``H_blocks`` prefix instead of the ``blocks`` prefix.
157 """
158 block_prefixes = ["L_blocks", "H_blocks"]
159 conversions: dict[str, ParamProcessingConversion | str] = {}
160 for prefix in block_prefixes:
161 conversions.update(
162 {
163 f"{prefix}.{{i}}.attn.q.weight": ParamProcessingConversion(
164 tensor_conversion=RearrangeTensorConversion(
165 "(n h) m -> n m h", n=self.cfg.n_heads
166 ),
167 ),
168 f"{prefix}.{{i}}.attn.k.weight": ParamProcessingConversion(
169 tensor_conversion=RearrangeTensorConversion(
170 "(n h) m -> n m h", n=n_kv_heads
171 ),
172 ),
173 f"{prefix}.{{i}}.attn.v.weight": ParamProcessingConversion(
174 tensor_conversion=RearrangeTensorConversion(
175 "(n h) m -> n m h", n=n_kv_heads
176 ),
177 ),
178 f"{prefix}.{{i}}.attn.o.weight": ParamProcessingConversion(
179 tensor_conversion=RearrangeTensorConversion(
180 "m (n h) -> n h m", n=self.cfg.n_heads
181 ),
182 ),
183 }
184 )
185 return conversions
187 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
188 """Set up rotary embedding references for HRM-Text component testing.
190 HRM-Text uses RoPE. We set the rotary_emb reference on all attention bridge
191 instances so component-level isolation tests can run.
192 """
193 rotary_emb = hf_model.model.rotary_emb
195 if hasattr(hf_model, "config") and hasattr(hf_model.config, "_attn_implementation"):
196 hf_model.config._attn_implementation = "eager"
198 for stack_attr in ("L_module", "H_module"):
199 stack = getattr(hf_model.model, stack_attr, None)
200 if stack is not None and hasattr(stack, "layers"):
201 for layer in stack.layers:
202 if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"):
203 layer.self_attn.config._attn_implementation = "eager"
205 if bridge_model is not None:
206 for blocks_attr in ("L_blocks", "H_blocks"):
207 blocks = getattr(bridge_model, blocks_attr, None)
208 if blocks is not None:
209 for block in blocks:
210 if hasattr(block, "attn"):
211 block.attn.set_rotary_emb(rotary_emb)
213 for blocks_path in ("L_blocks.0.attn", "H_blocks.0.attn"):
214 try:
215 attn_bridge = self.get_generalized_component(blocks_path)
216 attn_bridge.set_rotary_emb(rotary_emb)
217 except (KeyError, AttributeError):
218 pass