Add HRM-Text (hrm_text) hierarchical recurrent model with PrefixLM masking - #592
Open
justinchuby wants to merge 2 commits into
Open
Add HRM-Text (hrm_text) hierarchical recurrent model with PrefixLM masking#592justinchuby wants to merge 2 commits into
justinchuby wants to merge 2 commits into
Conversation
Adds production ONNX export for the HuggingFace `hrm_text` architecture
(`sapientinc/HRM-Text-1B`).
HRM-Text is not a flat decoder stack: it owns two independently-weighted
transformer stacks (`L_module`, `H_module`) and runs them in a fixed
recurrence, so a single forward performs `H_cycles * (L_cycles + 1)` stack
invocations. Upstream reflects that by inflating `num_hidden_layers` to
`num_layers_per_stack * H_cycles * (L_cycles + 1)` so its cache allocates one
slot per unique attention invocation; the exported graph uses the same
inflated count and emits slots in the same order, so `past_key_values.{i}` /
`present.{i}` line up 1:1 with HuggingFace.
Architecture specifics replicated:
- parameterless (scale-free) RMSNorm with fp32 variance accumulation
- sigmoid-gated attention output driven by a separate `gate_proj`
- token-embedding scaling by `1 / initializer_range`
- MHA (k/v sized by `num_attention_heads`), full causal attention
- no trailing model norm (each stack ends with its own `final_norm`)
`preprocess_weights` un-fuses the checkpoint's `attn.gqkv_proj`
(gate/q/k/v on dim 0) and `mlp.gate_up_proj`, matching upstream's
`conversion_mapping` entry, and stays an identity for already-converted
HuggingFace state dicts.
Supporting changes:
- `Attention._post_attention`: a no-op extension point invoked before
`o_proj` on both the plain and GroupQueryAttention paths, so gated
variants reuse the whole Q/K/V + RoPE + cache pipeline.
- `TextModel._build_attention_context`: the existing EP-aware GQA /
padding-mask / static-cache bias selection moved out of `forward` so a
model with a non-sequential layer schedule can reuse it verbatim.
- `load_torch_model` accepts `revision`, and `generate_golden.py` forwards
the test case revision so golden references are reproducibly pinned.
Tests: L1 tiny graph + weight alignment, L2 arch validation, L3 synthetic
parity, real-weight prefill and cached-decode parity against HuggingFace,
and dedicated unit tests for config inflation, cache-slot layout, and the
fused-weight split.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
HRM-Text was pre-trained with a PrefixLM mask: prompt tokens attend
bidirectionally within one prefix block, generated tokens attend causally.
The model card for `sapientinc/HRM-Text-1B` is explicit that omitting it
"does not match the pre-training distribution and will give noticeably
worse logits". The previous commit parsed `config.prefix_lm` but never
implemented it, so the export was unconditionally causal.
Measured on the pinned revision with HF fp32 on CPU, causal vs PrefixLM is
a real semantic difference rather than rounding:
prompt prefill logit delta (max / mean)
"Here is my poem:" 14.26 / 1.10
"The capital of France is" 16.48 / 1.28
model-card prompt 19.61 / 1.35
and only the model card's documented prompt under PrefixLM produces a
coherent, non-degenerate continuation.
Implementation:
- `CausalLMTask` gains an opt-in `token_type_ids` graph input, probed via a
`requires_token_type_ids` module hook exactly like the existing
`kv_cache_layer_count` / `static_kv_cache_specs` hooks. `task_type` stays
`text-generation` and no other model gains an input or forward keyword.
- `HrmTextModel` maps `block_sequence_ids = where(token_type_ids == 1, 0, -1)`
and feeds the existing `create_attention_bias(block_sequence_ids=...)`
primitive, which implements HF's `blockwise_overlay`
(`(q_group == kv_group) & (q_group >= 0)`, OR-ed onto causal) exactly.
- `Attention` gains an `_is_causal` instance attribute so a float additive
bias that already encodes causality is not double-masked. It is a static
per-graph property, so it needs no forward-signature change.
Upstream gates the overlay on `is_first_iteration`; the graph reproduces
that by data instead of by branching. A generated token is fed with
`token_type_ids == 0` -> block `-1`, and the `q_group >= 0` guard makes the
overlay a no-op, so decode stays causal. All-zero `token_type_ids` therefore
reproduces HF's `token_type_ids=None` path through the very same graph.
Also fixes a silent faithfulness bug found while testing: a raw
`config.json` only carries the default `rope_theta` of 10000.0, which the
generic extractor deliberately ignores as a RoPE signal, so the raw-config
path exported a position-free graph. `HrmTextRotaryEmbedding` is
unconditional upstream, so `HrmTextConfig` now pins default RoPE explicitly.
Golden/test plumbing sends the identical contract to both sides, derived
once in `mobius._testing.prefix_lm` rather than restated per call site:
`generate_golden.py`, `e2e_golden_test.py`, `OnnxGenerator`, and the
integration feeds. `torch_forward` accepts `token_type_ids` and forwards it
only when the model's signature has it.
The L4/L5 goldens are regenerated with the documented prompt and are now
non-degenerate (24 tokens, 23 distinct):
"The sky appears blue due to the scattering of sunlight by air molecules.
Shorter wavelengths of light (blue". The case sets `exact_match: true`, so
L5 asserts exact sequence length, token IDs, and decoded text.
Tests: PrefixLM prefill parity, causal-fallback parity through the same
graph, an assertion that the two modes diverge by the same amount on both
sides (so a graph ignoring the overlay cannot pass), PrefixLM cached decode,
graph-input declaration, and the `prefix_lm=False` input set.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Comment on lines
+291
to
+300
| def forward( | ||
| self, | ||
| op: OpBuilder, | ||
| input_ids: ir.Value, | ||
| attention_mask: ir.Value | None, | ||
| position_ids: ir.Value, | ||
| past_key_values: list | None = None, | ||
| inputs_embeds: ir.Value | None = None, | ||
| deepstack_embeds: list | None = None, | ||
| ): |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds first-class Mobius ONNX export support for the HuggingFace hrm_text architecture (HRM-Text-1B), including its hierarchical recurrent H/L stack schedule and PrefixLM masking semantics, and wires revision pinning through golden/reference generation and integration testing to keep results reproducible.
Changes:
- Add
HrmTextConfigextraction logic (including explicit default RoPE pinning) and register/export the newHrmTextCausalLMModel. - Implement HRM-Text model graph construction (H/L recurrence, gated attention output, embedding scaling) and opt-in
token_type_idshandling for PrefixLM viaCausalLMTask. - Add/update test + golden infrastructure for PrefixLM (
token_type_ids) and for pinned-revision golden/reference generation.
Reviewed changes
Copilot reviewed 20 out of 20 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/synthetic_parity_test.py | Adds hrm_text tiny-config overrides and feeds zero token_type_ids when the exported graph declares it. |
| tests/integration_test.py | Adds real-weight HRM-Text-1B PrefixLM prefill + cached decode parity test and threads token_type_ids into feeds. |
| tests/e2e_golden_test.py | Ensures e2e golden harness feeds correct token_type_ids contract for PrefixLM graphs. |
| tests/_test_configs.py | Adds a tiny HRM-Text config entry for graph-build and representative-arch coverage. |
| testdata/golden/causal-lm/hrm-text-1b.json | Adds L4 logits golden for HRM-Text-1B. |
| testdata/golden/causal-lm/hrm-text-1b_generation.json | Adds L5 greedy generation golden for HRM-Text-1B. |
| testdata/cases/causal-lm/hrm-text-1b.yaml | Adds an HRM-Text-1B test case including pinned revision and generation parameters. |
| src/mobius/tasks/_causal_lm.py | Adds optional token_type_ids graph input gated by module.requires_token_type_ids. |
| src/mobius/models/hrm_text.py | New HRM-Text backbone + recurrence implementation, PrefixLM bias path, and checkpoint weight unfusing. |
| src/mobius/models/hrm_text_test.py | Adds unit tests for config inflation, graph/cache shape, recurrence parity, and weight preprocessing. |
| src/mobius/models/base.py | Factors out shared attention-context construction into _build_attention_context for reuse by non-sequential schedulers. |
| src/mobius/models/init.py | Exports HrmTextCausalLMModel. |
| src/mobius/components/_attention.py | Adds _post_attention hook and per-module _is_causal control to support gated output and float-bias masking. |
| src/mobius/_testing/torch_reference.py | Adds revision pinning to load_torch_model and optional token_type_ids forwarding in torch_forward. |
| src/mobius/_testing/prefix_lm.py | Introduces a shared PrefixLM token_type_ids contract helper for tests/golden generation. |
| src/mobius/_testing/generation.py | Updates ONNX greedy generator to feed PrefixLM-style token_type_ids when required by the graph. |
| src/mobius/_registry.py | Registers hrm_text and adds a default test model id entry for the architecture. |
| src/mobius/_configs/_base.py | Adds HrmTextConfig dataclass and extraction logic (including explicit default RoPE pinning). |
| src/mobius/_configs/init.py | Exports HrmTextConfig. |
| scripts/generate_golden.py | Threads pinned revision into HF reference loading and applies PrefixLM token_type_ids contract for goldens. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+248
to
+265
| if token_type_ids is None: | ||
| raise ValueError( | ||
| "HRM-Text was built with config.prefix_lm=True, which requires a " | ||
| "token_type_ids input. The task supplies it via the module's " | ||
| "requires_token_type_ids hook; a caller invoking forward() " | ||
| "directly must pass token_type_ids explicitly." | ||
| ) | ||
| if attention_mask is None: | ||
| # Static cache passes attention_mask=None and relies on | ||
| # is_causal=1 + nonpad_kv_seqlen, which cannot express the | ||
| # bidirectional prefix block. | ||
| raise NotImplementedError( | ||
| "HRM-Text PrefixLM masking is not supported with the static KV " | ||
| "cache: the bidirectional prefix overlay needs the dynamic " | ||
| "attention_mask to build its float additive bias. Build with " | ||
| "the default dynamic cache, or set config.prefix_lm=False for a " | ||
| "fully causal export." | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds production ONNX export for the HuggingFace
hrm_textarchitecture, usingsapientinc/HRM-Text-1Bpinned to revision9f082d68b8cd0ebc56e33f1c88c45609174c272c. The revision is threaded through config, weights, tokenizer, golden generation, the integration test, and the CLI.Architecture
HRM-Text is not a flat decoder stack. It owns two independently-weighted transformer stacks — a fast "low" stack (
L_module) and a slow "high" stack (H_module) — and runs them in a fixed recurrence:Every stack invocation runs all
num_layers_per_stackblocks and therefore performs its own attention with its own KV-cache slots. Upstream reflects that by inflatingconfig.num_hidden_layerstonum_layers_per_stack * H_cycles * (L_cycles + 1)soDynamicCacheallocates one slot per unique attention invocation. The exported graph uses the same inflated count and emits slots in the same order, sopast_key_values.{i}/present.{i}line up 1:1 with HuggingFace. For HRM-Text-1B that is 128 cache slots driven by 32 blocks of weights (16 per stack), in a 2198-node graph.Faithfully replicated against
transformers.models.hrm_textandtransformers.conversion_mapping:HrmTextRMSNormhas no learnable scale and normalises in fp32 →ScaleFreeRMSNormgate_proj;sigmoid(gate) * attn_outapplied beforeo_proj, from the same normalized states that produced Q/K/V* config.embedding_scale, defaulting to1 / initializer_rangek_proj/v_projsized bynum_attention_heads(upstream hardcodesnum_key_value_groups = 1)final_norm, so the lastH_moduleoutput feedslm_headdirectlyblock_sequence_ids = where(token_type_ids == 1, 0, -1), bidirectional within the prefix blockattn.gqkv_projsplits gate, q, k, v on dim 0;mlp.gate_up_projsplits gate/upPrefixLM is implemented, not approximated
The model card is explicit that omitting
token_type_ids"does not match the pre-training distribution and will give noticeably worse logits". Measured on the pinned revision (HF fp32, CPU), this is a real semantic difference rather than rounding:"Here is my poem:":tokens"The capital of France is"So
CausalLMTaskgained an opt-intoken_type_idsgraph input, probed via arequires_token_type_idsmodule hook — the same pattern as the existingkv_cache_layer_count/static_kv_cache_specshooks.task_typestaystext-generation, the registry is untouched, and no other model gains an input or a forward keyword.HrmTextModelthen feeds mobius's existingcreate_attention_bias(block_sequence_ids=...)primitive, which implements HF'sblockwise_overlay—(q_group == kv_group) & (q_group >= 0), OR-ed onto causal before the padding AND — identically.Upstream gates the overlay on
is_first_iteration. The graph reproduces that by data rather than by branching: a generated token is fed withtoken_type_ids == 0→ block-1, and theq_group >= 0guard makes the overlay a pure no-op, so decode stays causal. Consequently all-zerotoken_type_idsreproduces HF'stoken_type_ids=Nonepath through the very same graph, and one export serves both contracts.Supporting changes to shared code
Attention._post_attention— a no-op extension point invoked beforeo_projon both the plain and theGroupQueryAttentionpath, so gated variants reuse the entire Q/K/V + RoPE + cache pipeline instead of duplicating it.Attention._is_causal— an instance attribute (default1) so a subclass feeding a float additive bias that already encodes causality is not double-masked. It is a static per-graph property, so it needs no forward-signature change. Static-cache mode still derives causality fromattn_maskpresence, unchanged.TextModel._build_attention_context— the existing EP-aware GQA / padding-mask / static-cache-bias selection moved out offorwardverbatim so a model with a non-sequential layer schedule can reuse it.mobius._testing.prefix_lm— single source of truth for the PrefixLM contract, so the HF reference and the ONNX graph cannot drift apart. Wired intogenerate_golden.py,e2e_golden_test.py,OnnxGenerator, and the integration feeds.load_torch_modelacceptsrevision, andgenerate_golden.pyforwards the test case revision, so golden references are reproducibly pinned.Bug found and fixed while testing
A raw
config.jsoncarries only the defaultrope_theta = 10000.0, which mobius's generic extractor deliberately ignores as a RoPE signal (NoPE models inherit it as dead config data). That leftrope_type = Noneand silently exported a position-free graph even thoughHrmTextRotaryEmbeddingis unconditional upstream.HrmTextConfignow pins default RoPE explicitly, with a regression test. Verified this cannot clobber a genuine non-default RoPE (rope_theta=500000andrope_scaling={"rope_type":"yarn"}both survive, because a non-default theta is itself a RoPE signal).Test evidence
All runs used
PYTHONPATH=<this worktree>/srcandPYTHONIOENCODING=utf-8, verified in-process viaprint(mobius.__file__)._test_configs.py,is_representative=True)atol=1e-3integration_test.py, pinned revision)exact_match: truehrm_textunit testsReal-weight numerical parity (
sapientinc/HRM-Text-1B, pinned)Full-logit comparison against an independently loaded HF fp32 CPU reference:
session.get_providers()was asserted in-script, not assumed; the run hard-fails if CUDA EP is absent.CPU-vs-CUDA: the identical fp16 graph produced byte-identical 24-token output on both EPs, and both equal the committed L5 golden exactly.
The overlay is provably live, not inert. Prefix-vs-causal divergence measured through the same exported graph: HF fp32 19.6093 · ONNX f16/cuda 19.6328 · bf16/cuda 19.6875 · f16/cpu 19.6211. Tests assert the ONNX gap matches HF's to 1% and that it exceeds
1.0on real weights, so a graph that silently dropped PrefixLM cannot pass.L5 golden (non-degenerate)
24 tokens, 23 distinct.
exact_match: trueasserts exact sequence length, exact token IDs, and exact decoded text.Graph shape
MatMul 1025, SkipSimplifiedLayerNorm 263, Mul 257, RotaryEmbedding 256, Attention 128, Sigmoid 128, Swish 128prefix_lm=False… GroupQueryAttention 128, Sigmoid 128, Swish 128128 attention nodes = exactly one per unique attention invocation, and 128
Sigmoid= the gated output, confirming the_post_attentionhook fires on both the plain and the fused GQA path. The PrefixLM graph forgoes GQA fusion because the float additive bias requiresis_causal=0and GQA's mask model is causal/local-window only — the same correctness-for-speed trade-off Gemma4 documents for its vision-block overlay.CLI + runtime metadata
mobius build --model sapientinc/HRM-Text-1B --revision 9f082d68… --runtime ort-genaisucceeds, producingmodel.onnx+model.onnx.data(4515.6 MB) +genai_config.json+ tokenizer files. The metadata is graph-faithful, includingnum_hidden_layers: 128(the inflated slot count, matching the real cache contract) and thetoken_type_idsinput.INT4 quantization
Quantizing the exported package with ORT's
MatMulNBitsQuantizer(4-bit, block 32, asymmetric) works end to end:MatMul→MatMulNBits, including the fusedgqkv_proj-derived projections andlm_head; the recurrence structure (Attention 128,Sigmoid 128,Swish 128,RotaryEmbedding 256) is preserved"The sky appears blue due to the scattering of sunlight by air molecules. Shorter wavelengths of light (blue"tests/quantization_integration_test.pyalso passes (6/6), confirming the sharedCausalLMTask/Attentionchanges introduce no quantization regressions.Waivers
1. onnxruntime-genai cannot drive this model (runtime-side, two independent reasons).
Probed with onnxruntime-genai 0.15.2 against the generated package:
I then re-probed a variant with
token_type_idsremoved from the config, to prove the failure is not caused by mobius emitting that input:So ORT GenAI rejects
hrm_textoutright regardless of PrefixLM — it has no registry entry for this architecture — and separately has notoken_type_idsvalue in its decoder input schema. Both are runtime limitations, not export defects. Per the quality checklist, downstream runtime capability does not gate Mobius export, and the generated metadata is kept faithful to the graph (a config omitting the required input would load and then fail later with a missing-input error, which is strictly worse).2. PrefixLM + static KV cache is not supported. The static-cache path passes
attention_mask=Noneand relies onis_causal=1+nonpad_kv_seqlen, which cannot express a bidirectional prefix block. This raises an explicitNotImplementedErrornaming the two supported alternatives, rather than silently exporting a degraded causal graph.3. Chunked-prefill edge case. Upstream gates the overlay on
is_first_iteration; the graph gates it per-query onblock_sequence_ids >= 0. These agree exactly for the documented single-shot prefill and for all decoding. They would differ only if a caller split one prefix block across multiple cached prefill chunks, which neither the model card norgenerate()does.Review
A specialist code review was run twice. Round 1's single High finding was the missing PrefixLM implementation — independently reaching the same conclusion as the analysis above — which is fixed here. Round 2 reported no significant issues, having independently verified overlay semantics against
masking_utils.blockwise_overlay, GQA suppression across cuda/f16, cpu/f32, dml/f16 and default/f32, that_is_causalis inert for all 7Attentionsubclasses across 121 model modules, that the RoPE pin cannot clobber genuine non-default RoPE, and that thebase.pyextraction is behavior-preserving.The only failing tests in the suite (
gguf q4_0,deepseek_v4QMoE,decompose_attentionsoftcap,qwen_imagecuda low-precision,modernbert-decodersynthetic parity) reproduce byte-identically on059bf5f1and are unrelated to this change.