Added HRM-Text two-timescale recurrent adapter#1487
Conversation
|
hey @jlarson4 PR is up, please review it |
|
@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 |
jlarson4
left a comment
There was a problem hiding this comment.
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.forwardnow returns a 3-tuple (was a bare Tensor. The adapter wraps the gate in a Tensor-typedLinearBridge. The fix: use the tuple-awareGeneralizedComponent(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.
|
hey @jlarson4, thanks for the comments. Please review it once again |
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:
BlockBridgeentries (L_blocks→model.L_module.layers,H_blocks→model.H_module.layers) with identical submodule structure (RMSNorm, PositionEmbeddingsAttentionBridge with sigmoid gate, GatedMLPBridge)supports_fold_ln = Falsesupports_center_writing_weights = Falseembedding_scalefolded intoembed.weightinpreprocess_weightsfor compat-mode parityH_cycles,L_cycles,L_bp_cycles,num_layers_per_stack,embedding_scale,prefix_lmRequired dependency bump:
transformers >= 5.9.0(HF introducedHrmTextForCausalLMin v5.9.0, currently at v5.12.1).Closes #1471
Type of change
Screenshots
N/A
Checklist: