Coverage for transformer_lens/model_bridge/supported_architectures/hrm_text.py: 55%
47 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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 LinearBridge,
50 PositionEmbeddingsAttentionBridge,
51 RMSNormalizationBridge,
52 RotaryEmbeddingBridge,
53 UnembeddingBridge,
54)
55from transformer_lens.utilities.attn_implementation import force_eager_attention
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._set_rms_rotary_defaults()
76 if hasattr(cfg, "num_key_value_heads") and cfg.num_key_value_heads is not None: 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true
77 self.cfg.n_key_value_heads = cfg.num_key_value_heads
78 elif hasattr(cfg, "num_attention_heads"): 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true
79 self.cfg.n_key_value_heads = cfg.num_attention_heads
81 for attr in (
82 "H_cycles",
83 "L_cycles",
84 "L_bp_cycles",
85 "num_layers_per_stack",
86 "embedding_scale",
87 "prefix_lm",
88 ):
89 if hasattr(cfg, attr):
90 setattr(self.cfg, attr, getattr(cfg, attr))
92 n_kv_heads = (
93 self.cfg.n_key_value_heads
94 if hasattr(self.cfg, "n_key_value_heads") and self.cfg.n_key_value_heads is not None
95 else self.cfg.n_heads
96 )
97 self.weight_processing_conversions = self._build_weight_conversions(n_kv_heads)
99 def _make_block_submodules():
100 return {
101 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
102 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
103 "attn": PositionEmbeddingsAttentionBridge(
104 name="self_attn",
105 config=self.cfg,
106 submodules={
107 "q": LinearBridge(name="q_proj"),
108 "k": LinearBridge(name="k_proj"),
109 "v": LinearBridge(name="v_proj"),
110 "o": LinearBridge(name="o_proj"),
111 "gate": LinearBridge(name="gate_proj"),
112 },
113 requires_attention_mask=True,
114 requires_position_embeddings=True,
115 ),
116 "mlp": self._gated_mlp(),
117 }
119 self.component_mapping = {
120 "embed": EmbeddingBridge(name="model.embed_tokens"),
121 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
122 "L_blocks": BlockBridge(
123 name="model.L_module.layers",
124 submodules=_make_block_submodules(),
125 ),
126 "H_blocks": BlockBridge(
127 name="model.H_module.layers",
128 submodules=_make_block_submodules(),
129 ),
130 "L_ln_final": RMSNormalizationBridge(name="model.L_module.final_norm", config=self.cfg),
131 "H_ln_final": RMSNormalizationBridge(name="model.H_module.final_norm", config=self.cfg),
132 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
133 }
135 def _build_weight_conversions(
136 self, n_kv_heads: int
137 ) -> dict[str, ParamProcessingConversion | str]:
138 """Build weight processing conversions for both L and H block stacks.
140 Each Q/K/V/O weight under ``L_blocks.{i}`` and ``H_blocks.{i}`` needs
141 the same ``(n_heads * d_head, d_model) → (n_heads, d_head, d_model)``
142 rearrangement as a standard decoder adapter, but with the ``L_blocks`` /
143 ``H_blocks`` prefix instead of the ``blocks`` prefix.
144 """
145 block_prefixes = ["L_blocks", "H_blocks"]
146 conversions: dict[str, ParamProcessingConversion | str] = {}
147 for prefix in block_prefixes:
148 conversions.update(
149 {
150 f"{prefix}.{{i}}.attn.q.weight": ParamProcessingConversion(
151 tensor_conversion=RearrangeTensorConversion(
152 "(n h) m -> n m h", n=self.cfg.n_heads
153 ),
154 ),
155 f"{prefix}.{{i}}.attn.k.weight": ParamProcessingConversion(
156 tensor_conversion=RearrangeTensorConversion(
157 "(n h) m -> n m h", n=n_kv_heads
158 ),
159 ),
160 f"{prefix}.{{i}}.attn.v.weight": ParamProcessingConversion(
161 tensor_conversion=RearrangeTensorConversion(
162 "(n h) m -> n m h", n=n_kv_heads
163 ),
164 ),
165 f"{prefix}.{{i}}.attn.o.weight": ParamProcessingConversion(
166 tensor_conversion=RearrangeTensorConversion(
167 "m (n h) -> n h m", n=self.cfg.n_heads
168 ),
169 ),
170 }
171 )
172 return conversions
174 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
175 """Set up rotary embedding references for HRM-Text component testing.
177 HRM-Text uses RoPE. We set the rotary_emb reference on all attention bridge
178 instances so component-level isolation tests can run.
179 """
180 rotary_emb = hf_model.model.rotary_emb
182 # per_layer covers both L_module and H_module stacks.
183 force_eager_attention(hf_model, per_layer=True)
185 if bridge_model is not None:
186 for blocks_attr in ("L_blocks", "H_blocks"):
187 blocks = getattr(bridge_model, blocks_attr, None)
188 if blocks is not None:
189 for block in blocks:
190 if hasattr(block, "attn"):
191 block.attn.set_rotary_emb(rotary_emb)
193 for blocks_path in ("L_blocks.0.attn", "H_blocks.0.attn"):
194 try:
195 attn_bridge = self.get_generalized_component(blocks_path)
196 attn_bridge.set_rotary_emb(rotary_emb)
197 except (KeyError, AttributeError):
198 pass