transformer_lens.model_bridge.supported_architectures.hrm_text module

HRM-Text architecture adapter.

HRM-Text (Sapient Intelligence) is a hierarchical two-timescale recurrent model: two transformer stacks (H = slow/planning, L = fast/computation) iterate in a nested loop with additive cross-stack coupling.

Architecture notes:

  • Two physical stacks: model.L_module and model.H_module, each with num_layers_per_stack layers. The stacks share identical internal structure but have separate weights.

  • Recurrence: outer H-cycle iterates H_cycles times; each iteration runs L_cycles inner L-cycle iterations. Total forward passes through the layer stacks = H_cycles * (L_cycles + 1). The config field num_hidden_layers is rewritten by HF to num_layers_per_stack * H_cycles * (L_cycles + 1) to size the KV cache slots.

  • Parameterless RMSNorm: input_layernorm, post_attention_layernorm, and each stack’s final_norm have no learnable weight tensor.

  • Sigmoid attention gate: each attention block has a gate_proj linear that produces a per-head sigmoid gate applied to the attention output before o_proj. Delegated to HF; hookable via L_blocks.{i}.attn.gate.hook_out.

  • Embedding scale: inputs_embeds *= embedding_scale (default ~39.19 for HRM-Text-1B). Applied at runtime by HrmTextModel.forward; must NOT be folded into embed.weight — same reasoning as gemma1.py.

  • PrefixLM mask: instruction tokens attend bidirectionally when token_type_ids is passed to HF forward; delegated, not modeled by bridge.

Known limitations:

  1. Hooks on L_blocks.{i}.* fire H_cycles * L_cycles times per forward; on H_blocks.{i}.* they fire H_cycles times. No per-iteration index is exposed; per-cycle disambiguation is future work.

  2. Compat-mode with PrefixLM (token_type_ids) inputs is untested in v1.

  3. supports_fold_ln = False — parameterless norms cannot be folded.

  4. supports_center_writing_weights = False — block naming (L_blocks / H_blocks instead of blocks) is incompatible with weight-centering iteration over range(cfg.n_layers).

  5. Requires transformers >= 5.9.0 at runtime.

class transformer_lens.model_bridge.supported_architectures.hrm_text.HrmTextArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for HRM-Text (Sapient Intelligence).

Exposes L_blocks (fast/low-level stack) and H_blocks (slow/high-level stack) as sibling block lists. The nested recurrence loop is owned by HF’s forward; hooks fire once per iteration through the physical layers.

__init__(cfg: Any) → None

Initialize the HRM-Text architecture adapter.

applicable_phases: list[int] = [1, 2, 3]
component_mapping: ComponentMapping | None
setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up rotary embedding references for HRM-Text component testing.

HRM-Text uses RoPE. We set the rotary_emb reference on all attention bridge instances so component-level isolation tests can run.

supports_center_writing_weights = False
supports_fold_ln = False
uses_split_attention: bool
weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None