Skip to content

Added HRM-Text two-timescale recurrent adapter#1487

Open
MdSadiqMd wants to merge 12 commits into
TransformerLensOrg:devfrom
MdSadiqMd:sadiq/HrmTextForCausalLM-adapter
Open

Added HRM-Text two-timescale recurrent adapter#1487
MdSadiqMd wants to merge 12 commits into
TransformerLensOrg:devfrom
MdSadiqMd:sadiq/HrmTextForCausalLM-adapter

Conversation

@MdSadiqMd

Copy link
Copy Markdown
Contributor

Description

Add a TransformerBridge adapter for HrmTextForCausalLM (Sapient Intelligence HRM-Text), a hierarchical two-timescale recurrent model. The adapter exposes the fast (L) and slow (H) stacks as side-by-side hookable block lists (L_blocks, H_blocks) so the two timescales can be ablated and patched independently.

HRM-Text's nested recurrence (H_cycles × (L_cycles + 1) forward passes through the reused physical stacks) is a third flavor of depth-recurrence alongside Huginn's variable-count loop and Ouro's halting loop. The dual-stack architecture gives interp researchers a clean substrate to study fixed-point attractor dynamics, cross-timescale state coupling, and latent-space reasoning — none of which are expressible on a flat block list.

Key adapter properties:

  • Two BlockBridge entries (L_blocksmodel.L_module.layers, H_blocksmodel.H_module.layers) with identical submodule structure (RMSNorm, PositionEmbeddingsAttentionBridge with sigmoid gate, GatedMLPBridge)
  • Parameterless RMSNorm → supports_fold_ln = False
  • Non-standard block names → supports_center_writing_weights = False
  • embedding_scale folded into embed.weight in preprocess_weights for compat-mode parity
  • Passthrough attrs: H_cycles, L_cycles, L_bp_cycles, num_layers_per_stack, embedding_scale, prefix_lm

Required dependency bump: transformers >= 5.9.0 (HF introduced HrmTextForCausalLM in v5.9.0, currently at v5.12.1).

Closes #1471

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

Screenshots

N/A

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@MdSadiqMd

Copy link
Copy Markdown
Contributor Author

hey @jlarson4 PR is up, please review it

@jlarson4 jlarson4 changed the base branch from main to dev July 7, 2026 04:04
@jlarson4

jlarson4 commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

@MdSadiqMd It'll be a few hours before I can do a full review of this, if you want to work on it in the meantime, it appears your CI run has a couple test failures. I'll make sure to ping you when I've done a more thorough review, thank you for putting this together!

Also, I resolved conflicts with the latest version of dev, I'd recommend pulling those changes down before making any changes.

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @MdSadiqMd! Great progress on this so far, you've done a great job handling all the tricky bits that come with this architecture. There is a bit of work yet to be done before we can merge this, however. I have added a few specific comments as well as some general points below.

1. Booted bridge crashes on default forward and generate() (KV-cache / cycle_offset)

Booting through the production path (TransformerBridge.boot_transformers on a saved tiny HrmText checkpoint) succeeds, but a plain bridge(ids) and bridge.generate() both crash: RuntimeError: The size of tensor a (16) must match the size of tensor b (8) at non-singleton dimension 3. HF's HrmText attention updates the KV cache at self.layer_idx + cycle_offset (modeling_hrm_text.py:215) so each of the H_cycles×(L_cycles+1) stack invocations gets its own cache slot; the bridge's reimplemented attention updates at self._layer_idx only (generalized_components/attention.py:370-379, ignores the cycle_offset kwarg), so repeated invocations of the same physical layer within one forward append K/V into one slot. An 8-token prompt yields 16 cached keys and a mask-shape mismatch.

When use_cache=False the bridge is exactly correct, which is why none of the PR's tests catch this: the integration tests bypass boot() entirely and hardcode use_cache=False on every call. You'll need to fix the cache slotting (respect cycle_offset).

2. Compat mode breaks parity

The embedding_scale fold double-applies preprocess_weights folds embedding_scale (×8 on the tiny config) into embed.weight, and compat mode writes the folded weight back into the live HF module, but HF's HrmTextModel.forward still multiplies inputs_embeds * self.embedding_scale at runtime (modeling_hrm_text.py:467), so the scale applies twice. gemma1.py:51-56 documents this trap ("must NOT pre-scale weights" when HF applies the scale at runtime). The fix: drop the fold, using the Gemma example as a point of reference.

3. The transformers 5.13 adaptation needs a little more work

Commit d2d2b4a ("pipeline fix") legitimately fixed the unit-tier MLA failure and the unit tier is now fully green, But the Full Code Coverage job fails 4 integration tests, all transformers-5.13 fallout.

  • glm_moe_dsa parity: max diff 0.0016867, bit-identical across runs
    HF 5.13 switched GLM-MoE-DSA main-attention RoPE from split-half NeoX style to interleaved-pair rotation (apply_rotary_pos_emb_interleave). The bridge's _apply_rotary_pos_emb_single (glm_moe_dsa_attention.py:19-24, applied at :138/:140) still implements the 5.8.1 form.
  • glm4_moe: beartype tuple violation. 5.13's Glm4MoeTopkRouter.forward now returns a 3-tuple (was a bare Tensor. The adapter wraps the gate in a Tensor-typed LinearBridge. The fix: use the tuple-aware GeneralizedComponent(name="gate", optional=True). deepseek_v3.py:82 and glm_moe_dsa.py:75 are good examples of how this works.
  • granite_moe_hybrid + nemotron_h layers_block_type: HF 5.13 canonically remaps legacy layer
    type names at config construction ('mamba' maps to 'linear_attention', 'attention' maps to 'full_attention', configuration_utils.py:82-92). The adapters are pure passthroughs and nothing in TL branches on the strings (SSM detection is structural), so the right fix is updating the two tests' expected lists to the canonical names plus the stale legacy vocabulary mentioned in nemotron_h.py's docstring.

Latent 5.13 gaps in the same component, unexercised by tests:
With attention_mask=None the bridge's DSA index mask lacks HF's causal constraint (future-token leakage in standalone use), and a shared layer with prev_topk_indices=None calls a None indexer (TypeError) where HF raises a clean ValueError.

4. Missing a "verify_models run" for a HRM-Text model

Add a sapientinc/HRM-Text-1B entry to tools/model_registry/data/supported_models.json and run verify_models --model sapientinc/HRM-Text-1B. A verification run will catch a lot of the latent bugs that have not yet been addressed.

5. L_blocks/H_blocks silently miss every blocks-hardcoded bridge feature The adapter spec says the block container is "always named blocks", but this is an instance that proves this false. The T5 adapter's encoder_blocks/decoder_blocks are the only special-cased exceptions that are currently in place. The fix: implement adapter-level setup_hook_compatibility over L_blocks/H_blocks (the bridge already has an adapter branch for this) so q/k/v/z hooks get head-split shapes, and make stop_at_layer raise instead of silently no-op.

Longer-term, I will need to make an issue for us to generalize the bridge's block-list discovery from the literal blocks attribute to all BlockBridge lists in component_mapping.

Comment thread tests/integration/model_bridge/test_hrm_text_adapter.py
Comment thread transformer_lens/model_bridge/supported_architectures/hrm_text.py Outdated
@MdSadiqMd

Copy link
Copy Markdown
Contributor Author

hey @jlarson4, thanks for the comments. Please review it once again

Comment thread transformer_lens/model_bridge/architecture_adapter.py Outdated
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Proposal] Add HRM-Text two-timescale recurrent adapter (HrmTextForCausalLM)

2 participants