From 1c7a46a2c87d5c8fecf2af1488eaf1c18f2ed1cc Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 24 Aug 2026 05:52:00 -0700 Subject: [PATCH 1/2] Add HRM-Text (hrm_text) hierarchical recurrent model support 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 --- scripts/generate_golden.py | 5 +- src/mobius/_configs/__init__.py | 2 + src/mobius/_configs/_base.py | 90 ++++++ src/mobius/_registry.py | 4 + src/mobius/_testing/torch_reference.py | 9 +- src/mobius/components/_attention.py | 30 ++ src/mobius/models/__init__.py | 2 + src/mobius/models/base.py | 136 +++++--- src/mobius/models/hrm_text.py | 377 ++++++++++++++++++++++ src/mobius/models/hrm_text_test.py | 372 +++++++++++++++++++++ testdata/cases/causal-lm/hrm-text-1b.yaml | 23 ++ tests/_test_configs.py | 18 ++ tests/integration_test.py | 98 ++++++ tests/synthetic_parity_test.py | 8 + 14 files changed, 1119 insertions(+), 55 deletions(-) create mode 100644 src/mobius/models/hrm_text.py create mode 100644 src/mobius/models/hrm_text_test.py create mode 100644 testdata/cases/causal-lm/hrm-text-1b.yaml diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index e05617ba7..e31d6acb6 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -267,7 +267,10 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: ) else: model, tokenizer = load_torch_model( - case.model_id, device=device, trust_remote_code=case.trust_remote_code + case.model_id, + device=device, + trust_remote_code=case.trust_remote_code, + revision=case.revision, ) encoded = tokenizer(case.prompts[0], return_tensors="np", padding=False) diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index 2fbc36c8f..bc0444e40 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -38,6 +38,7 @@ Gemma4Config, GlmAsrConfig, GraniteMoeHybridConfig, + HrmTextConfig, JambaConfig, JetMoeConfig, Lfm2Config, @@ -114,6 +115,7 @@ "Gemma4Config", "GlmAsrConfig", "GraniteMoeHybridConfig", + "HrmTextConfig", "JambaConfig", "JetMoeConfig", "Lfm2Config", diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index cafe91f3c..5355cc873 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -1739,6 +1739,96 @@ def from_transformers(cls, config, parent_config=None) -> Gemma2Config: ) +def _as_int_or_default(value, default: int) -> int: + """Return ``int(value)``, falling back to *default* only when unset. + + Unlike ``int(value or default)`` this preserves an explicit ``0`` so that + an out-of-range checkpoint value reaches the config's own validation + instead of being silently rewritten to the default. + """ + return default if value is None else int(value) + + +def _as_float_or_default(value, default: float) -> float: + """Float counterpart of :func:`_as_int_or_default`.""" + return default if value is None else float(value) + + +@dataclasses.dataclass +class HrmTextConfig(CausalLMConfig): + """Configuration for HRM-Text hierarchical recurrent models. + + Mirrors HuggingFace ``HrmTextConfig``. The checkpoint stores the real + per-stack block count in ``num_hidden_layers``; HuggingFace's + ``__post_init__`` moves it to ``num_layers_per_stack`` and rewrites + ``num_hidden_layers`` to the *inflated* total number of unique attention + invocations under the H/L recurrence:: + + num_hidden_layers = num_layers_per_stack * H_cycles * (L_cycles + 1) + + That inflated count is what drives KV-cache slot allocation, so it is what + :class:`~mobius.tasks.CausalLMTask` must see. This dataclass reproduces the + same split so that both a trusted ``HrmTextConfig`` instance (already + inflated, ``num_layers_per_stack`` set) and a raw pinned ``config.json`` + (not inflated, ``num_layers_per_stack`` absent) resolve identically. + + ``embedding_scale`` defaults to ``1 / initializer_range`` exactly as + upstream does when the checkpoint leaves it unset. + """ + + H_cycles: int = 2 + L_cycles: int = 3 + num_layers_per_stack: int | None = None + embedding_scale: float | None = None + initializer_range: float = 0.02 + prefix_lm: bool = True + + def __post_init__(self): + if self.H_cycles <= 0 or self.L_cycles <= 0: + raise ValueError( + f"HrmTextConfig requires positive H_cycles/L_cycles, got " + f"H_cycles={self.H_cycles}, L_cycles={self.L_cycles}" + ) + # HRM-Text attention is always MHA: upstream hardcodes + # ``num_key_value_groups = 1`` and sizes k_proj/v_proj by + # ``num_attention_heads * head_dim``, ignoring any + # ``num_key_value_heads`` the checkpoint happens to carry. + if self.num_attention_heads != DEFAULT_INT: + self.num_key_value_heads = self.num_attention_heads + if self.embedding_scale is None: + if not self.initializer_range: + raise ValueError( + "HrmTextConfig needs a non-zero initializer_range to derive " + "embedding_scale when the checkpoint does not supply one." + ) + self.embedding_scale = 1.0 / self.initializer_range + if self.num_layers_per_stack is None and self.num_hidden_layers != DEFAULT_INT: + # Raw-config path: ``num_hidden_layers`` still carries the real + # per-stack count. Remember it, then inflate exactly as upstream. + self.num_layers_per_stack = self.num_hidden_layers + self.num_hidden_layers = ( + self.num_layers_per_stack * self.H_cycles * (self.L_cycles + 1) + ) + + @classmethod + def from_transformers(cls, config, parent_config=None) -> HrmTextConfig: + base = ArchitectureConfig.from_transformers(config, parent_config) + fields = _shallow_fields(base) + return cls( + **fields, + H_cycles=_as_int_or_default(getattr(config, "H_cycles", None), 2), + L_cycles=_as_int_or_default(getattr(config, "L_cycles", None), 3), + # Present only on a trusted HF config that already ran its own + # ``__post_init__``; ``None`` for a raw ``config.json``. + num_layers_per_stack=getattr(config, "num_layers_per_stack", None), + embedding_scale=getattr(config, "embedding_scale", None), + initializer_range=_as_float_or_default( + getattr(config, "initializer_range", None), 0.02 + ), + prefix_lm=bool(getattr(config, "prefix_lm", True)), + ) + + @dataclasses.dataclass class NanoChatConfig(CausalLMConfig): """Configuration for NanoChat models with final logit soft-capping. diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index d626c1f7b..d395b940f 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -28,6 +28,7 @@ Gemma3nMultiModalConfig, Gemma4AssistantConfig, Gemma4Config, + HrmTextConfig, Lfm2Config, Lfm2VlConfig, MMSConfig, @@ -73,6 +74,7 @@ GPTOSSCausalLMModel, GraniteCausalLMModel, GraniteMoECausalLMModel, + HrmTextCausalLMModel, HunYuanMoEV1CausalLMModel, HunYuanV1DenseCausalLMModel, HunYuanVLMoTModel, @@ -492,6 +494,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "gpt_oss": ModelRegistration(GPTOSSCausalLMModel), "gptj": ModelRegistration(GPTJCausalLMModel), "granite": ModelRegistration(GraniteCausalLMModel), + "hrm_text": ModelRegistration(HrmTextCausalLMModel, config_class=HrmTextConfig), "hunyuan_v1_dense": ModelRegistration(HunYuanV1DenseCausalLMModel), "internlm2": ModelRegistration(InternLM2CausalLMModel), "llama4_text": ModelRegistration(Llama4CausalLMModel), @@ -953,6 +956,7 @@ def _create_default_registry() -> ModelRegistry: "dots1": "rednote-hilab/dots.llm1.inst", "exaone4": "LGAI-EXAONE/EXAONE-4.0-1.2B", "helium": "kyutai/helium-1-preview-2b", + "hrm_text": "sapientinc/HRM-Text-1B", "minicpm": "optimum-intel-internal-testing/tiny-random-minicpm", "minicpm3": "openbmb/MiniCPM3-4B", "ministral3": "Aratako/Ministral-3-3B-Instruct-2512-BF16-TextOnly", diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 271ec2693..e3f71783e 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -198,6 +198,7 @@ def load_torch_model( dtype: torch.dtype = torch.float32, device: str = "cpu", trust_remote_code: bool = True, + revision: str | None = None, ): """Load a HuggingFace causal LM model for reference inference. @@ -211,6 +212,9 @@ def load_torch_model( natively supported by the installed transformers, so the transformers-5.x-compatible implementation is used instead of an older bundled ``modeling_*.py`` that relies on removed cache APIs. + revision: Hub revision to pin. ``None`` resolves the default branch; + golden-data generation always passes the test case's revision so + the reference and the exported graph read the same commit. Returns: Tuple of (model, tokenizer). @@ -220,14 +224,14 @@ def load_torch_model( _install_dynamic_cache_legacy_shims() tokenizer = transformers.AutoTokenizer.from_pretrained( - model_id, trust_remote_code=trust_remote_code + model_id, trust_remote_code=trust_remote_code, revision=revision ) # NemotronH: disable rescale_prenorm_residual before loading to # prevent _init_weights from corrupting out_proj.weight with # random kaiming_uniform_ initialization after checkpoint loading. config = transformers.AutoConfig.from_pretrained( - model_id, trust_remote_code=trust_remote_code + model_id, trust_remote_code=trust_remote_code, revision=revision ) if getattr(config, "model_type", None) == "nemotron_h": config.rescale_prenorm_residual = False @@ -238,6 +242,7 @@ def load_torch_model( dtype=dtype, device_map=device, trust_remote_code=trust_remote_code, + revision=revision, ) _fix_nemotron_h_init_weights(model, model_id) model.eval() diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index f6724e841..0c2f3309a 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -349,6 +349,7 @@ def forward( value_states, attention_bias, past_key_value, + hidden_states=hidden_states, ) # Apply rotary position embeddings (skip when not provided) @@ -396,9 +397,32 @@ def forward( static_cache=static_cache, ) + attn_output = self._post_attention(op, attn_output, hidden_states) attn_output = self.o_proj(op, attn_output) return attn_output, (present_key, present_value) + def _post_attention( + self, + op: OpBuilder, + attn_output: ir.Value, + hidden_states: ir.Value, + ) -> ir.Value: + """Transform the attention output before the ``o_proj`` projection. + + Extension point for architectures that post-process the attended + values while still reusing the whole Q/K/V + RoPE + cache pipeline of + this class. The base implementation is the identity, so it is inert + for every standard model. + + Args: + attn_output: Attention result ``[B, S, num_heads * head_dim]``. + hidden_states: The (already normalized) layer input that produced + Q/K/V, so a subclass can derive a gate from the same tensor + HuggingFace does. + """ + del op, hidden_states + return attn_output + def _forward_gqa( self, op: OpBuilder, @@ -407,6 +431,7 @@ def _forward_gqa( value_states: ir.Value, gqa_ctx: GQAContext, past_key_value: tuple | None, + hidden_states: ir.Value | None = None, ): """Emit ``com.microsoft::GroupQueryAttention`` directly. @@ -417,6 +442,9 @@ def _forward_gqa( :class:`~mobius.rewrite_rules._group_query_attention.RotaryAttentionToGQA` rewrite rule; RoPE is handled by the ``do_rotary=1`` attribute instead. + ``hidden_states`` is only forwarded to :meth:`_post_attention`; it is + optional so that existing positional callers keep working. + Returns ``(attn_output, (present_key, present_value))`` in the same shape as the standard :meth:`forward` path. """ @@ -455,6 +483,8 @@ def _forward_gqa( **gqa_attrs, ) + if hidden_states is not None: + attn_out = self._post_attention(op, attn_out, hidden_states) attn_out = self.o_proj(op, attn_out) return attn_out, (present_key, present_value) diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 57c617fac..ddd8c3964 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -70,6 +70,7 @@ "GraniteCausalLMModel", "GraniteMoECausalLMModel", "GraniteMoeHybridCausalLMModel", + "HrmTextCausalLMModel", "HunYuanMoEV1CausalLMModel", "HunYuanV1DenseCausalLMModel", "HunYuanVLMoTModel", @@ -244,6 +245,7 @@ from mobius.models.gptoss import GPTOSSCausalLMModel from mobius.models.granite import GraniteCausalLMModel, GraniteMoECausalLMModel from mobius.models.granitemoehybrid import GraniteMoeHybridCausalLMModel +from mobius.models.hrm_text import HrmTextCausalLMModel from mobius.models.hunyuan_dit import HunyuanDiT2DModel from mobius.models.hunyuan_v1 import HunYuanV1DenseCausalLMModel from mobius.models.hunyuan_vl_mot import HunYuanVLMoTModel diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py index 3819e9218..8a3a1bef3 100644 --- a/src/mobius/models/base.py +++ b/src/mobius/models/base.py @@ -164,21 +164,30 @@ def _maybe_static_cache_bias( dtype=self._dtype, ) - def forward( + def _build_attention_context( self, op: OpBuilder, - input_ids: ir.Value, + *, + input_ids: ir.Value | None, 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, - ): - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_tokens(op, input_ids) + hidden_states: ir.Value, + past_key_values: list | None, + ) -> tuple[GQAContext | ir.Value | None, tuple | None]: + """Build the per-graph attention bias and RoPE embeddings. + Returns ``(attention_bias, position_embeddings)``. Both are shared by + every decoder layer in the graph, so this runs once per forward. + Split out of :meth:`forward` so that models with a non-sequential layer + schedule (e.g. the HRM-Text H/L recurrence, which invokes its two + stacks many times over one shared mask) can reuse the exact same + EP-aware GQA / padding-mask / static-cache decision logic. + + Args: + input_ids: ``[B, S]`` token ids, or ``None`` for embeds-driven + forwards (the padding mask then keys off ``hidden_states``). + hidden_states: ``[B, S, hidden]`` embeddings for this step. + """ # Determine whether to emit GroupQueryAttention directly. # Conditions: # - attention_mask present: static-cache mode passes None; GQA requires seqlens_k. @@ -235,50 +244,73 @@ def forward( # position_embeddings not needed: GroupQueryAttention handles RoPE # internally via do_rotary=1. Passing None skips apply_rotary_pos_emb # in Attention.forward() (which checks `if position_embeddings is not None`). + return attention_bias, None + + # This path (CPU fp32, DML, non-fused RoPE, mRoPE, static cache) + # builds at most a bool padding mask; it has no way to express a + # sliding window. Warn if the model expects one so the divergence + # from HuggingFace for sequences longer than the window is not + # silent. (For seq <= window the result is identical regardless.) + if self._gqa_local_window_size() > 0: + logger.warning( + "Model declares a uniform sliding window " + "(sliding_window=%s) but is being built through a non-GQA " + "attention path (build dtype=%s); the exported graph uses " + "full causal attention and will diverge from HuggingFace " + "for sequences longer than the window. Build with a " + "GQA-capable execution provider/dtype (e.g. CUDA or DML " + "with float16/bfloat16) to apply the window.", + getattr(self.config, "sliding_window", None), + dtype, + ) + # NoPE models (e.g. NemotronH, GraniteMoeHybrid) have + # ``rotary_emb = None`` because ``initialize_rope`` returned + # ``None`` for ``config.rope_type is None``. Skip building + # position_embeddings so that Attention.forward sees + # ``position_embeddings=None`` and does not apply rotary encoding. + if self.rotary_emb is not None: + position_embeddings = self.rotary_emb(op, position_ids) + else: position_embeddings = None + + # When attention_mask is None (static cache mode), skip mask + # creation entirely — the Attention op uses is_causal=1 instead. + # When present, create a bool padding mask. Causal masking is + # handled by is_causal=1 on the Attention op (set in + # _apply_attention), so we only need padding information here. + if attention_mask is not None: + attention_bias = create_padding_mask( + op, + input_ids=hidden_states if input_ids is None else input_ids, + attention_mask=attention_mask, + ) else: - # This path (CPU fp32, DML, non-fused RoPE, mRoPE, static cache) - # builds at most a bool padding mask; it has no way to express a - # sliding window. Warn if the model expects one so the divergence - # from HuggingFace for sequences longer than the window is not - # silent. (For seq <= window the result is identical regardless.) - if self._gqa_local_window_size() > 0: - logger.warning( - "Model declares a uniform sliding window " - "(sliding_window=%s) but is being built through a non-GQA " - "attention path (build dtype=%s); the exported graph uses " - "full causal attention and will diverge from HuggingFace " - "for sequences longer than the window. Build with a " - "GQA-capable execution provider/dtype (e.g. CUDA or DML " - "with float16/bfloat16) to apply the window.", - getattr(self.config, "sliding_window", None), - dtype, - ) - # NoPE models (e.g. NemotronH, GraniteMoeHybrid) have - # ``rotary_emb = None`` because ``initialize_rope`` returned - # ``None`` for ``config.rope_type is None``. Skip building - # position_embeddings so that Attention.forward sees - # ``position_embeddings=None`` and does not apply rotary encoding. - if self.rotary_emb is not None: - position_embeddings = self.rotary_emb(op, position_ids) - else: - position_embeddings = None - - # When attention_mask is None (static cache mode), skip mask - # creation entirely — the Attention op uses is_causal=1 instead. - # When present, create a bool padding mask. Causal masking is - # handled by is_causal=1 on the Attention op (set in - # _apply_attention), so we only need padding information here. - if attention_mask is not None: - attention_bias = create_padding_mask( - op, - input_ids=hidden_states if input_ids is None else input_ids, - attention_mask=attention_mask, - ) - else: - attention_bias = self._maybe_static_cache_bias( - op, hidden_states, past_key_values - ) + attention_bias = self._maybe_static_cache_bias(op, hidden_states, past_key_values) + return attention_bias, position_embeddings + + 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, + ): + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(op, input_ids) + + attention_bias, position_embeddings = self._build_attention_context( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + hidden_states=hidden_states, + past_key_values=past_key_values, + ) present_key_values = [] output_layer_indices = getattr(self, "output_layer_indices", None) diff --git a/src/mobius/models/hrm_text.py b/src/mobius/models/hrm_text.py new file mode 100644 index 000000000..e4c014eb5 --- /dev/null +++ b/src/mobius/models/hrm_text.py @@ -0,0 +1,377 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""HRM-Text hierarchical recurrent text models (``sapientinc/HRM-Text-1B``). + +Replicates HuggingFace's ``HrmTextForCausalLM``. HRM-Text is *not* a plain +stack of decoder layers: 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:: + + z_H = embed(input_ids) * embedding_scale + z_L = z_L_init # broadcast over (B, S, H) + for h in range(H_cycles): + for l in range(L_cycles): + z_L = L_module(z_L + z_H) + z_H = H_module(z_H + z_L) + logits = lm_head(z_H) + +Every stack invocation runs all ``num_layers_per_stack`` blocks and therefore +performs its own attention with its own KV-cache slots. Upstream reflects that +by inflating ``config.num_hidden_layers`` to +``num_layers_per_stack * H_cycles * (L_cycles + 1)`` so ``DynamicCache`` +allocates one slot per unique attention invocation; the exported ONNX graph +uses the same inflated count and emits the slots in the same order, so the +graph's ``past_key_values.{i}`` / ``present.{i}`` indices line up 1:1 with +HuggingFace's cache layout. + +Architectural differences from the standard :class:`CausalLMModel`: + +* **Parameterless RMSNorm** — ``HrmTextRMSNorm`` has no learnable scale and + normalises in float32, matching :class:`ScaleFreeRMSNorm`. +* **Gated attention output** — an extra ``gate_proj`` produces a per-element + sigmoid gate applied to the attention result before ``o_proj``. +* **Embedding scaling** — token embeddings are multiplied by + ``config.embedding_scale`` (``1 / initializer_range`` when unset). +* **MHA, not GQA** — ``k_proj``/``v_proj`` are sized by ``num_attention_heads``. +* **No trailing model norm** — each stack ends with its own ``final_norm``, so + the last ``H_module`` output feeds ``lm_head`` directly. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig, HrmTextConfig +from mobius._weight_utils import split_gate_up_proj +from mobius.components import ( + Attention, + DecoderLayer, + Embedding, + Linear, + ScaleFreeRMSNorm, + initialize_rope, +) +from mobius.models.base import CausalLMModel, TextModel + +if TYPE_CHECKING: + import onnx_ir as ir + +# Order of the four equally-sized chunks packed into the checkpoint's fused +# ``attn.gqkv_proj`` weight, matching HuggingFace's ``hrm_text`` entry in +# ``transformers.conversion_mapping`` (``Chunk(dim=0)`` over gate/q/k/v). +_FUSED_GQKV_PARTS: tuple[str, ...] = ("gate_proj", "q_proj", "k_proj", "v_proj") + +_FUSED_GQKV_SUFFIX = ".attn.gqkv_proj.weight" +_FUSED_GATE_UP_SUFFIX = ".mlp.gate_up_proj.weight" + + +class HrmTextAttention(Attention): + """Multi-head attention with a sigmoid output gate. + + Identical to the base :class:`~mobius.components.Attention` except for an + extra ``gate_proj`` that is driven by the *same* (already normalised) + layer input as Q/K/V. Upstream applies the gate on the per-head + ``(B, S, num_heads, head_dim)`` view; because both tensors use that exact + layout, the elementwise product is identical on the flattened + ``(B, S, num_heads * head_dim)`` view used here. + """ + + def __init__(self, config: ArchitectureConfig, linear_class: type | None = None): + super().__init__(config, linear_class=linear_class) + gate_linear = linear_class if linear_class is not None else Linear + self.gate_proj = gate_linear( + self.hidden_size, + self.num_attention_heads * self.head_dim, + bias=config.attn_qkv_bias, + ) + + def _post_attention( + self, + op: OpBuilder, + attn_output: ir.Value, + hidden_states: ir.Value, + ) -> ir.Value: + # gate: (B, S, num_heads * head_dim); attn_output has the same shape. + gate = self.gate_proj(op, hidden_states) + return op.Mul(attn_output, op.Sigmoid(gate)) + + +class HrmTextDecoderLayer(DecoderLayer): + """Pre-norm decoder layer with parameterless norms and gated attention. + + Replicates HuggingFace's ``HrmTextDecoderLayer``: ``input_layernorm`` → + gated self-attention → residual → ``post_attention_layernorm`` → SwiGLU + MLP → residual. Both norms are scale-free RMSNorms, so the checkpoint + ships no weights for them. + """ + + def __init__(self, config: ArchitectureConfig): + super().__init__(config, norm_class=ScaleFreeRMSNorm) + # Replace the plain attention built by DecoderLayer with the gated + # variant; the discarded module's parameters are dropped with it. + self.self_attn = HrmTextAttention(config) + + +class HrmTextStack(nn.Module): + """One HRM transformer stack (used twice: as ``L_module`` and ``H_module``). + + Replicates HuggingFace's ``HrmTextStack``: ``num_layers_per_stack`` + decoder layers followed by a parameterless ``final_norm``. The same + instance is invoked several times per forward pass, once per recurrence + step, so its parameters appear exactly once in the ONNX graph while its + ops are unrolled per invocation. + """ + + def __init__(self, config: ArchitectureConfig, num_layers: int): + super().__init__() + self.layers = nn.ModuleList([HrmTextDecoderLayer(config) for _ in range(num_layers)]) + self.final_norm = ScaleFreeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias, + position_embeddings: tuple | None, + past_key_values: list, + ): + """Run every block in the stack over one recurrence step. + + Args: + past_key_values: Exactly ``len(self.layers)`` cache entries — the + slice of the global cache belonging to *this* invocation. + + Returns: + ``(hidden_states, present_key_values)`` where ``hidden_states`` is + ``(B, S, hidden)`` after ``final_norm``. + """ + if len(past_key_values) != len(self.layers): + raise ValueError( + f"HrmTextStack expected {len(self.layers)} cache entries for this " + f"invocation, got {len(past_key_values)}" + ) + present_key_values = [] + for layer, past_kv in zip(self.layers, past_key_values): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + return self.final_norm(op, hidden_states), present_key_values + + +class HrmTextModel(TextModel): + """HRM-Text backbone: embeddings + the H/L recurrence over two stacks. + + Replicates HuggingFace's ``HrmTextModel``. Subclasses + :class:`~mobius.models.base.TextModel` purely to reuse its EP-aware + attention-context construction (GQA fusion vs. bool padding mask vs. + static-cache bias); the layer schedule itself is completely different, so + ``forward`` is overridden and ``self.layers`` is intentionally absent. + """ + + def __init__(self, config: ArchitectureConfig): + # Deliberately skip TextModel.__init__: it would build a flat + # ``self.layers`` stack and a trailing ``self.norm`` that HRM-Text + # does not have. + nn.Module.__init__(self) + self.config = config + self._dtype = config.dtype + self.output_layer_indices = None + # HRM-Text uses full causal attention; no sliding window. Required by + # the inherited ``_maybe_static_cache_bias`` / ``_gqa_local_window_size``. + self._sliding_window = None + + self._h_cycles = int(config.H_cycles) + self._l_cycles = int(config.L_cycles) + self._layers_per_stack = _resolve_layers_per_stack(config) + self._embedding_scale = float(config.embedding_scale) + + self.embed_tokens = Embedding( + config.vocab_size, config.hidden_size, config.pad_token_id + ) + self.rotary_emb = initialize_rope(config) + self.L_module = HrmTextStack(config, self._layers_per_stack) + self.H_module = HrmTextStack(config, self._layers_per_stack) + # Frozen initial low-cycle state, stored as a (hidden,) vector and + # broadcast against (B, S, hidden) — equivalent to HF's ``expand_as``. + self.z_L_init = nn.Parameter([config.hidden_size], dtype=config.dtype) + + 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, + ): + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(op, input_ids) + # Token-embedding multiplier (1 / initializer_range for this family). + # z_H — slow / high-level state: (B, S, hidden) + z_high = op.Mul(hidden_states, self._embedding_scale) + + attention_bias, position_embeddings = self._build_attention_context( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + hidden_states=z_high, + past_key_values=past_key_values, + ) + + num_invocations = self._h_cycles * (self._l_cycles + 1) + expected_slots = num_invocations * self._layers_per_stack + cache = past_key_values if past_key_values is not None else [None] * expected_slots + if len(cache) != expected_slots: + raise ValueError( + f"HRM-Text expects {expected_slots} KV-cache slots " + f"({num_invocations} stack invocations x {self._layers_per_stack} " + f"layers), got {len(cache)}" + ) + + # z_L — fast / low-level state. Starts as the frozen (hidden,) vector + # and is broadcast by the first Add against z_high. + z_low: ir.Value = self.z_L_init + present_key_values: list = [] + # Cache slots are consumed strictly in invocation order, which + # reproduces upstream's + # slot(h, l, layer) = (h * (L_cycles + 1) + l) * layers_per_stack + layer + # (the trailing H invocation of cycle h uses l == L_cycles). + slot = 0 + for _ in range(self._h_cycles): + for _ in range(self._l_cycles): + next_slot = slot + self._layers_per_stack + z_low, presents = self.L_module( + op, + hidden_states=op.Add(z_low, z_high), + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_values=cache[slot:next_slot], + ) + present_key_values.extend(presents) + slot = next_slot + + next_slot = slot + self._layers_per_stack + z_high, presents = self.H_module( + op, + hidden_states=op.Add(z_high, z_low), + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_values=cache[slot:next_slot], + ) + present_key_values.extend(presents) + slot = next_slot + + # No trailing model-level norm: H_module.final_norm already ran. + return z_high, present_key_values + + +class HrmTextCausalLMModel(CausalLMModel): + """Causal LM head over the HRM-Text hierarchical recurrent backbone. + + Replicates HuggingFace's ``HrmTextForCausalLM`` (``sapientinc/HRM-Text-1B``): + two recurrently-invoked transformer stacks with parameterless RMSNorm, + sigmoid-gated attention output, and scaled token embeddings. + + Inputs: ``input_ids``, ``attention_mask``, ``position_ids``, + ``past_key_values``. Outputs: ``logits`` and one present KV pair per + unique attention invocation of the recurrence. + """ + + config_class: type = HrmTextConfig + + def __init__(self, config: ArchitectureConfig): + nn.Module.__init__(self) + self.config = config + self.model = HrmTextModel(config) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + if config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Map HRM-Text checkpoint keys onto this module tree. + + Handles both layouts that reach us: + + * The **raw checkpoint** (``sapientinc/HRM-Text-1B``), which packs + ``attn.gqkv_proj`` (gate/q/k/v concatenated on dim 0) and + ``mlp.gate_up_proj``, and names the attention submodule ``attn``. + * An **already-converted** HuggingFace ``state_dict``, whose keys equal + our ONNX parameter names — that path is a pure identity mapping. + """ + head_width = self.config.num_attention_heads * self.config.head_dim + intermediate_size = self.config.intermediate_size + converted: dict[str, torch.Tensor] = {} + + for key, value in state_dict.items(): + if key.endswith(_FUSED_GQKV_SUFFIX): + prefix = key[: -len(_FUSED_GQKV_SUFFIX)] + expected = len(_FUSED_GQKV_PARTS) * head_width + if value.shape[0] != expected: + raise ValueError( + f"{key}: fused gqkv_proj dim 0 is {value.shape[0]}, expected " + f"{expected} ({len(_FUSED_GQKV_PARTS)} x " + f"num_attention_heads * head_dim = {head_width})" + ) + for index, part in enumerate(_FUSED_GQKV_PARTS): + start = index * head_width + converted[f"{prefix}.self_attn.{part}.weight"] = value[ + start : start + head_width + ] + continue + + if key.endswith(_FUSED_GATE_UP_SUFFIX): + prefix = key[: -len(_FUSED_GATE_UP_SUFFIX)] + gate, up = split_gate_up_proj(value, intermediate_size) + converted[f"{prefix}.mlp.gate_proj.weight"] = gate + converted[f"{prefix}.mlp.up_proj.weight"] = up + continue + + # ``.attn.`` never matches the already-converted ``.self_attn.`` + # spelling, so this rename is idempotent. + converted[key.replace(".attn.", ".self_attn.")] = value + + return super().preprocess_weights(converted) + + +def _resolve_layers_per_stack(config: ArchitectureConfig) -> int: + """Return the number of transformer blocks inside each H / L stack. + + ``config.num_hidden_layers`` is the *inflated* per-invocation count, so the + real per-stack depth comes from ``num_layers_per_stack`` when the config + carries it and is otherwise derived back out of the inflated total. + """ + per_stack = getattr(config, "num_layers_per_stack", None) + h_cycles = int(config.H_cycles) + l_cycles = int(config.L_cycles) + invocations = h_cycles * (l_cycles + 1) + if per_stack is None: + per_stack, remainder = divmod(int(config.num_hidden_layers), invocations) + if remainder or per_stack <= 0: + raise ValueError( + f"HRM-Text num_hidden_layers ({config.num_hidden_layers}) is not a " + f"positive multiple of H_cycles * (L_cycles + 1) = {invocations}" + ) + return per_stack + + per_stack = int(per_stack) + expected_total = per_stack * invocations + if int(config.num_hidden_layers) != expected_total: + raise ValueError( + f"HRM-Text num_hidden_layers ({config.num_hidden_layers}) must equal " + f"num_layers_per_stack * H_cycles * (L_cycles + 1) = {expected_total}" + ) + return per_stack diff --git a/src/mobius/models/hrm_text_test.py b/src/mobius/models/hrm_text_test.py new file mode 100644 index 000000000..836eb87b3 --- /dev/null +++ b/src/mobius/models/hrm_text_test.py @@ -0,0 +1,372 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the HRM-Text hierarchical recurrent export. + +Covers the three things that are unique to this architecture and that a +graph-build test cannot reach: + +* config extraction — ``num_hidden_layers`` is inflated to one slot per unique + attention invocation, from *both* a trusted HuggingFace config (already + inflated) and a raw pinned ``config.json`` (not inflated); +* ``preprocess_weights`` — the checkpoint packs gate/q/k/v into + ``attn.gqkv_proj`` and gate/up into ``mlp.gate_up_proj``; +* the H/L recurrence itself — prefill *and* a cached decode step are compared + against HuggingFace, which is what actually pins the KV-cache slot order. + +All tiny random configs -- no checkpoint download. +""" + +from __future__ import annotations + +import types + +import numpy as np +import onnx_ir as ir +import pytest +import torch + +from mobius._configs import HrmTextConfig +from mobius._testing.ort_inference import OnnxModelSession +from mobius.integrations._weight_loading import apply_weights +from mobius.models.hrm_text import ( + HrmTextCausalLMModel, + _resolve_layers_per_stack, +) +from mobius.tasks import get_task + +_HIDDEN = 64 +_HEADS = 4 +_HEAD_DIM = 16 +_INTERMEDIATE = 128 +_VOCAB = 256 +_PER_STACK = 2 +_H_CYCLES = 2 +_L_CYCLES = 2 +# One cache slot per unique attention invocation. +_TOTAL_SLOTS = _PER_STACK * _H_CYCLES * (_L_CYCLES + 1) + + +def _raw_json_config(**overrides) -> types.SimpleNamespace: + """A stand-in for a raw pinned ``config.json`` (no HF ``__post_init__``).""" + fields = { + "model_type": "hrm_text", + "vocab_size": _VOCAB, + "hidden_size": _HIDDEN, + "intermediate_size": _INTERMEDIATE, + # Raw checkpoints carry the *per-stack* depth here. + "num_hidden_layers": _PER_STACK, + "num_attention_heads": _HEADS, + "num_key_value_heads": _HEADS, + "head_dim": _HEAD_DIM, + "H_cycles": _H_CYCLES, + "L_cycles": _L_CYCLES, + "max_position_embeddings": 128, + "rms_norm_eps": 1e-6, + "rope_theta": 10_000.0, + "tie_word_embeddings": False, + "initializer_range": 0.025, + "prefix_lm": True, + "pad_token_id": 0, + "hidden_act": "silu", + } + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +def _mobius_config(**overrides) -> HrmTextConfig: + config = HrmTextConfig.from_transformers(_raw_json_config(**overrides)) + config.dtype = ir.DataType.FLOAT + return config + + +def _hf_config(**overrides): + """Build the upstream ``HrmTextConfig`` for the same tiny architecture.""" + transformers = pytest.importorskip("transformers") + fields = { + "hidden_size": _HIDDEN, + "intermediate_size": _INTERMEDIATE, + "num_attention_heads": _HEADS, + "num_hidden_layers": _PER_STACK, + "vocab_size": _VOCAB, + "max_position_embeddings": 128, + "rms_norm_eps": 1e-6, + "pad_token_id": 0, + "head_dim": _HEAD_DIM, + "H_cycles": _H_CYCLES, + "L_cycles": _L_CYCLES, + "initializer_range": 0.025, + "prefix_lm": True, + "rope_parameters": {"rope_type": "default", "rope_theta": 10_000.0}, + } + fields.update(overrides) + return transformers.AutoConfig.for_model("hrm_text", **fields) + + +# --------------------------------------------------------------------------- +# Config extraction +# --------------------------------------------------------------------------- + + +def test_raw_json_config_inflates_layer_count(): + config = _mobius_config() + assert config.num_layers_per_stack == _PER_STACK + assert config.num_hidden_layers == _TOTAL_SLOTS + + +def test_trusted_hf_config_keeps_inflated_layer_count(): + hf_config = _hf_config() + # Upstream already inflated in its own __post_init__. + assert hf_config.num_hidden_layers == _TOTAL_SLOTS + assert hf_config.num_layers_per_stack == _PER_STACK + + config = HrmTextConfig.from_transformers(hf_config) + assert config.num_hidden_layers == _TOTAL_SLOTS + assert config.num_layers_per_stack == _PER_STACK + + +def test_config_forces_multi_head_attention(): + # HrmTextAttention hardcodes num_key_value_groups = 1, so a checkpoint + # claiming grouped-query heads must not shrink k_proj/v_proj. + config = _mobius_config(num_key_value_heads=1) + assert config.num_key_value_heads == config.num_attention_heads == _HEADS + + +def test_embedding_scale_defaults_to_inverse_initializer_range(): + raw = _raw_json_config(initializer_range=0.025) + del raw.prefix_lm + config = HrmTextConfig.from_transformers(raw) + assert config.embedding_scale == pytest.approx(1.0 / 0.025) + # ``prefix_lm`` defaults to True upstream. + assert config.prefix_lm is True + + +def test_explicit_embedding_scale_is_preserved(): + config = _mobius_config(embedding_scale=39.191835884530846) + assert config.embedding_scale == pytest.approx(39.191835884530846) + + +def test_config_rejects_non_positive_cycles(): + with pytest.raises(ValueError, match="positive H_cycles"): + HrmTextConfig.from_transformers(_raw_json_config(L_cycles=0)) + + +def test_resolve_layers_per_stack_rejects_inconsistent_total(): + config = _mobius_config() + config.num_hidden_layers = _TOTAL_SLOTS + 1 + with pytest.raises(ValueError, match="num_layers_per_stack"): + _resolve_layers_per_stack(config) + + +def test_resolve_layers_per_stack_derives_from_total_when_unset(): + config = _mobius_config() + config.num_layers_per_stack = None + assert _resolve_layers_per_stack(config) == _PER_STACK + + +# --------------------------------------------------------------------------- +# Graph shape +# --------------------------------------------------------------------------- + + +def _build_package(config: HrmTextConfig): + module = HrmTextCausalLMModel(config) + return module, get_task("text-generation").build(module, config) + + +def test_graph_exposes_one_cache_slot_per_attention_invocation(): + config = _mobius_config() + _, pkg = _build_package(config) + graph = pkg["model"].graph + input_names = {value.name for value in graph.inputs} + output_names = {value.name for value in graph.outputs} + for slot in range(_TOTAL_SLOTS): + assert f"past_key_values.{slot}.key" in input_names + assert f"past_key_values.{slot}.value" in input_names + assert f"present.{slot}.key" in output_names + assert f"present.{slot}.value" in output_names + assert f"past_key_values.{_TOTAL_SLOTS}.key" not in input_names + + +def test_stack_weights_are_shared_across_recurrence_steps(): + config = _mobius_config() + _, pkg = _build_package(config) + names = set(pkg["model"].graph.initializers) + # Two stacks x _PER_STACK layers of parameters, regardless of how many + # times the recurrence invokes them. + for stack in ("L_module", "H_module"): + for layer in range(_PER_STACK): + prefix = f"model.{stack}.layers.{layer}" + assert f"{prefix}.self_attn.q_proj.weight" in names + assert f"{prefix}.self_attn.gate_proj.weight" in names + assert f"model.{stack}.layers.{_PER_STACK}.self_attn.q_proj.weight" not in names + assert "model.z_L_init" in names + + +# --------------------------------------------------------------------------- +# Weight preprocessing +# --------------------------------------------------------------------------- + + +def _fused_checkpoint_state_dict(config: HrmTextConfig) -> dict[str, torch.Tensor]: + """Mimic the layout of ``sapientinc/HRM-Text-1B``'s ``model.safetensors``.""" + head_width = config.num_attention_heads * config.head_dim + state: dict[str, torch.Tensor] = { + "model.embed_tokens.weight": torch.randn(config.vocab_size, config.hidden_size), + "model.z_L_init": torch.zeros(config.hidden_size), + "lm_head.weight": torch.randn(config.vocab_size, config.hidden_size), + } + for stack in ("L_module", "H_module"): + for layer in range(config.num_layers_per_stack): + prefix = f"model.{stack}.layers.{layer}" + state[f"{prefix}.attn.gqkv_proj.weight"] = torch.randn( + 4 * head_width, config.hidden_size + ) + state[f"{prefix}.attn.o_proj.weight"] = torch.randn(config.hidden_size, head_width) + state[f"{prefix}.mlp.gate_up_proj.weight"] = torch.randn( + 2 * config.intermediate_size, config.hidden_size + ) + state[f"{prefix}.mlp.down_proj.weight"] = torch.randn( + config.hidden_size, config.intermediate_size + ) + return state + + +def test_preprocess_weights_unfuses_checkpoint_layout(): + config = _mobius_config() + module = HrmTextCausalLMModel(config) + state = _fused_checkpoint_state_dict(config) + result = module.preprocess_weights(state) + + head_width = config.num_attention_heads * config.head_dim + prefix = "model.L_module.layers.0" + fused = state[f"{prefix}.attn.gqkv_proj.weight"] + # Upstream conversion_mapping order: gate, q, k, v. + for index, part in enumerate(("gate_proj", "q_proj", "k_proj", "v_proj")): + expected = fused[index * head_width : (index + 1) * head_width] + torch.testing.assert_close(result[f"{prefix}.self_attn.{part}.weight"], expected) + + fused_mlp = state[f"{prefix}.mlp.gate_up_proj.weight"] + torch.testing.assert_close( + result[f"{prefix}.mlp.gate_proj.weight"], fused_mlp[: config.intermediate_size] + ) + torch.testing.assert_close( + result[f"{prefix}.mlp.up_proj.weight"], fused_mlp[config.intermediate_size :] + ) + # ``attn`` -> ``self_attn`` rename, and no fused key survives. + assert f"{prefix}.self_attn.o_proj.weight" in result + assert not any("gqkv_proj" in key or "gate_up_proj" in key for key in result) + + +def test_preprocess_weights_covers_every_graph_parameter(): + config = _mobius_config() + module, pkg = _build_package(config) + result = module.preprocess_weights(_fused_checkpoint_state_dict(config)) + missing = { + name + for name, init in pkg["model"].graph.initializers.items() + if init.const_value is None and name not in result + } + assert not missing, sorted(missing) + + +def test_preprocess_weights_is_identity_for_converted_names(): + config = _mobius_config() + module, pkg = _build_package(config) + aligned = { + name: torch.ones(list(init.shape)) + for name, init in pkg["model"].graph.initializers.items() + if init.const_value is None + } + result = module.preprocess_weights(aligned) + assert set(aligned) <= set(result) + + +def test_preprocess_weights_rejects_wrong_fused_width(): + config = _mobius_config() + module = HrmTextCausalLMModel(config) + state = _fused_checkpoint_state_dict(config) + state["model.L_module.layers.0.attn.gqkv_proj.weight"] = torch.randn( + 3 * config.num_attention_heads * config.head_dim, config.hidden_size + ) + with pytest.raises(ValueError, match="fused gqkv_proj"): + module.preprocess_weights(state) + + +# --------------------------------------------------------------------------- +# Numerical parity against HuggingFace (prefill + cached decode) +# --------------------------------------------------------------------------- + + +def test_recurrence_matches_huggingface_prefill_and_decode(): + """Prefill *and* a cached decode step must match upstream. + + The decode step is the part that pins the KV-cache slot layout: an + off-by-one stack ordering still passes a cacheless prefill comparison but + reads the wrong slots on the second step. + """ + transformers = pytest.importorskip("transformers") + torch.manual_seed(11) + hf_config = _hf_config() + hf_model = transformers.AutoModelForCausalLM.from_config(hf_config).float().eval() + + config = HrmTextConfig.from_transformers(hf_config) + config.dtype = ir.DataType.FLOAT + module, pkg = _build_package(config) + apply_weights(pkg["model"], module.preprocess_weights(dict(hf_model.state_dict()))) + + rng = np.random.default_rng(11) + input_ids = rng.integers(1, _VOCAB, size=(1, 5)).astype(np.int64) + attention_mask = np.ones_like(input_ids) + position_ids = np.arange(input_ids.shape[1], dtype=np.int64)[np.newaxis, :] + + session = OnnxModelSession(pkg["model"]) + try: + feeds: dict[str, np.ndarray] = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + } + empty = np.zeros((1, _HEADS, 0, _HEAD_DIM), dtype=np.float32) + for slot in range(_TOTAL_SLOTS): + feeds[f"past_key_values.{slot}.key"] = empty + feeds[f"past_key_values.{slot}.value"] = empty + prefill = session.run(feeds) + + with torch.no_grad(): + hf_prefill = hf_model( + input_ids=torch.from_numpy(input_ids), + attention_mask=torch.from_numpy(attention_mask), + position_ids=torch.from_numpy(position_ids), + use_cache=True, + ) + np.testing.assert_allclose( + prefill["logits"], hf_prefill.logits.numpy(), rtol=1e-3, atol=1e-3 + ) + + next_id = np.array([[int(hf_prefill.logits[0, -1].argmax())]], dtype=np.int64) + decode_mask = np.ones((1, input_ids.shape[1] + 1), dtype=np.int64) + decode_pos = np.array([[input_ids.shape[1]]], dtype=np.int64) + decode_feeds: dict[str, np.ndarray] = { + "input_ids": next_id, + "attention_mask": decode_mask, + "position_ids": decode_pos, + } + for slot in range(_TOTAL_SLOTS): + decode_feeds[f"past_key_values.{slot}.key"] = prefill[f"present.{slot}.key"] + decode_feeds[f"past_key_values.{slot}.value"] = prefill[f"present.{slot}.value"] + decode = session.run(decode_feeds) + + with torch.no_grad(): + hf_decode = hf_model( + input_ids=torch.from_numpy(next_id), + attention_mask=torch.from_numpy(decode_mask), + position_ids=torch.from_numpy(decode_pos), + past_key_values=hf_prefill.past_key_values, + use_cache=True, + ) + np.testing.assert_allclose( + decode["logits"], hf_decode.logits.numpy(), rtol=1e-3, atol=1e-3 + ) + finally: + session.close() diff --git a/testdata/cases/causal-lm/hrm-text-1b.yaml b/testdata/cases/causal-lm/hrm-text-1b.yaml new file mode 100644 index 000000000..69e11e0cc --- /dev/null +++ b/testdata/cases/causal-lm/hrm-text-1b.yaml @@ -0,0 +1,23 @@ +model_id: "sapientinc/HRM-Text-1B" +model_type: "hrm_text" +revision: "9f082d68b8cd0ebc56e33f1c88c45609174c272c" +task_type: "text-generation" +dtype: "float32" + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 24 + do_sample: false + +notes: >- + HRM-Text 1B hierarchical recurrent decoder. Two independently-weighted + transformer stacks (L_module / H_module) are invoked H_cycles * (L_cycles + 1) + = 8 times per forward, so the graph carries 128 KV-cache slots for 16 + per-stack blocks. Parameterless RMSNorm, sigmoid-gated attention output, and + token embeddings scaled by 1 / initializer_range. Base model with no chat + template, so the prompt is fed through the plain tokenizer. diff --git a/tests/_test_configs.py b/tests/_test_configs.py index cbcc7ad6b..cf5dbb4bb 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -31,6 +31,7 @@ Gemma4Config, GlmAsrConfig, GraniteMoeHybridConfig, + HrmTextConfig, JambaConfig, JetMoeConfig, Lfm2Config, @@ -189,6 +190,23 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: False, ), ("helium", {}, False), + ( + # HRM-Text: two recurrently-invoked stacks. ``num_hidden_layers`` here + # is the real per-stack depth; ``HrmTextConfig.__post_init__`` inflates + # it to num_layers_per_stack * H_cycles * (L_cycles + 1) = 2 * 2 * 3 = 12 + # exactly as HuggingFace's ``HrmTextConfig`` does, so the tiny HF + # reference in the L3 parity test sees the same layer schedule. + "hrm_text", + { + "_config_cls": HrmTextConfig, + "num_hidden_layers": 2, + "H_cycles": 2, + "L_cycles": 2, + "initializer_range": 0.02, + "prefix_lm": True, + }, + True, + ), ("hunyuan_v1_dense", {}, False), ("llama4_text", {}, False), ("ministral", {}, False), diff --git a/tests/integration_test.py b/tests/integration_test.py index 0c535dd28..233b97df0 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -244,6 +244,104 @@ def test_minicpmv4_6_real_weight_vision_parity(): np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=1e-3) +_HRM_TEXT_MODEL_ID = "sapientinc/HRM-Text-1B" +_HRM_TEXT_REVISION = "9f082d68b8cd0ebc56e33f1c88c45609174c272c" + + +@pytest.mark.integration +@pytest.mark.integration_slow +def test_hrm_text_1b_real_weight_prefill_and_decode_parity(): + """Real HRM-Text-1B weights: prefill and cached decode must match HF. + + The cached decode step is the load-bearing half: HRM-Text drives 128 + KV-cache slots from only 32 blocks of weights, so a slot-ordering mistake + in the H/L recurrence still passes a prefill-only comparison but reads the + wrong cache entries on the second step. + + ``sapientinc/HRM-Text-1B`` is a base model with no chat template, so the + prompt goes through the plain tokenizer. The pinned revision is forwarded + to the config, tokenizer, weights, and the mobius build alike. + """ + import gc + + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + from mobius._configs import HrmTextConfig + + hf_config = AutoConfig.from_pretrained(_HRM_TEXT_MODEL_ID, revision=_HRM_TEXT_REVISION) + config = HrmTextConfig.from_transformers(hf_config) + # Upstream inflates num_hidden_layers to one slot per unique attention + # invocation; the exported graph must expose exactly that many. + assert config.num_layers_per_stack is not None + assert config.num_hidden_layers == ( + config.num_layers_per_stack * config.H_cycles * (config.L_cycles + 1) + ) + + tokenizer = AutoTokenizer.from_pretrained(_HRM_TEXT_MODEL_ID, revision=_HRM_TEXT_REVISION) + tokens = tokenizer("Here is my poem:", return_tensors="np") + input_ids = tokens["input_ids"].astype(np.int64) + attention_mask = tokens["attention_mask"].astype(np.int64) + seq_len = input_ids.shape[1] + position_ids = np.arange(seq_len, dtype=np.int64)[np.newaxis, :] + + torch_model = AutoModelForCausalLM.from_pretrained( + _HRM_TEXT_MODEL_ID, + revision=_HRM_TEXT_REVISION, + dtype=torch.float32, + ).eval() + with torch.no_grad(): + hf_prefill = torch_model( + input_ids=torch.from_numpy(input_ids), + attention_mask=torch.from_numpy(attention_mask), + position_ids=torch.from_numpy(position_ids), + use_cache=True, + ) + hf_prefill_logits = hf_prefill.logits.float().numpy() + + next_token = np.array([[int(hf_prefill_logits[0, -1].argmax())]], dtype=np.int64) + decode_mask = np.ones((1, seq_len + 1), dtype=np.int64) + decode_position_ids = np.array([[seq_len]], dtype=np.int64) + with torch.no_grad(): + hf_decode_logits = ( + torch_model( + input_ids=torch.from_numpy(next_token), + attention_mask=torch.from_numpy(decode_mask), + position_ids=torch.from_numpy(decode_position_ids), + past_key_values=hf_prefill.past_key_values, + use_cache=True, + ) + .logits.float() + .numpy() + ) + del torch_model, hf_prefill + gc.collect() + + package = build( + _HRM_TEXT_MODEL_ID, + dtype="f32", + load_weights=True, + revision=_HRM_TEXT_REVISION, + ) + session = _make_session(package["model"]) + try: + cache_inputs = [ + name for name in session.input_names if name.startswith("past_key_values.") + ] + assert len(cache_inputs) == 2 * config.num_hidden_layers + + feeds = _make_prefill_feeds(config, input_ids, attention_mask, position_ids) + onnx_prefill = session.run(feeds) + assert_logits_close(onnx_prefill["logits"], hf_prefill_logits, rtol=1e-3, atol=1e-3) + + decode_feeds = _make_decode_feeds( + config, next_token, decode_mask, decode_position_ids, onnx_prefill + ) + onnx_decode = session.run(decode_feeds) + assert_logits_close(onnx_decode["logits"], hf_decode_logits, rtol=1e-3, atol=1e-3) + finally: + session.close() + + # --------------------------------------------------------------------------- # Model catalogue: small models for each supported architecture # diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index 3b60e623a..ef874c158 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -319,6 +319,14 @@ "ministral3": {"head_dim": TINY_HEAD_DIM}, # Helium defaults head_dim=None in HF (causes pow(None,float) error) "helium": {"head_dim": TINY_HEAD_DIM}, + # HRM-Text: head_dim is an explicit config param (HF default 128), and + # rope settings live under rope_parameters rather than rope_theta. + # num_hidden_layers stays at the tiny *per-stack* depth so that HF's + # HrmTextConfig.__post_init__ inflates it the same way ours does. + "hrm_text": { + "head_dim": TINY_HEAD_DIM, + "rope_parameters": {"rope_type": "default", "rope_theta": 10_000.0}, + }, # seed_oss defaults head_dim=128 in HF; override to match tiny config "seed_oss": {"head_dim": TINY_HEAD_DIM}, # HunYuan V1 dense defaults head_dim=None in HF (causes pow(None,float) error) From 993050db1a6002a6c916ef6ce6f6a17ffa984cfb Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 24 Aug 2026 06:30:42 -0700 Subject: [PATCH 2/2] Implement HRM-Text PrefixLM masking via an opt-in token_type_ids input 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 --- scripts/generate_golden.py | 35 ++++- src/mobius/_configs/_base.py | 13 ++ src/mobius/_testing/generation.py | 17 ++ src/mobius/_testing/prefix_lm.py | 54 +++++++ src/mobius/_testing/torch_reference.py | 8 + src/mobius/components/_attention.py | 8 + src/mobius/models/hrm_text.py | 145 +++++++++++++++-- src/mobius/models/hrm_text_test.py | 148 +++++++++++++++--- src/mobius/tasks/_causal_lm.py | 18 +++ testdata/cases/causal-lm/hrm-text-1b.yaml | 9 +- testdata/golden/causal-lm/hrm-text-1b.json | 50 ++++++ .../causal-lm/hrm-text-1b_generation.json | 31 ++++ tests/e2e_golden_test.py | 16 +- tests/integration_test.py | 63 +++++++- tests/synthetic_parity_test.py | 6 + 15 files changed, 571 insertions(+), 50 deletions(-) create mode 100644 src/mobius/_testing/prefix_lm.py create mode 100644 testdata/golden/causal-lm/hrm-text-1b.json create mode 100644 testdata/golden/causal-lm/hrm-text-1b_generation.json diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index e31d6acb6..f0c17ce07 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -244,6 +244,10 @@ def _get_model_device(model: object, device: str): def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: """Generate golden data for a causal-lm (text-generation) model.""" from mobius._testing.golden import save_generation_json, save_golden_ref + from mobius._testing.prefix_lm import ( + model_type_uses_prefix_lm, + prompt_token_type_ids, + ) from mobius._testing.torch_reference import ( load_torch_model, load_torch_multimodal_model, @@ -279,6 +283,13 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: seq_len = input_ids.shape[1] position_ids = np.arange(seq_len).reshape(1, -1) + # PrefixLM models (HRM-Text) must run the reference with the same + # ``token_type_ids`` contract the exported graph is fed: the whole prompt is + # one bidirectional prefix block. Omitting it silently falls back to causal + # masking, which the model card warns does not match pre-training. + uses_prefix_lm = model_type_uses_prefix_lm(case.model_type) + prefix_token_type_ids = prompt_token_type_ids(input_ids) if uses_prefix_lm else None + # L4: single forward pass → last-token logits if uses_multimodal_reference: with torch.no_grad(): @@ -290,7 +301,13 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: ) logits = outputs.logits.float().cpu().numpy() else: - logits, _ = torch_forward(model, input_ids, attention_mask, position_ids) + logits, _ = torch_forward( + model, + input_ids, + attention_mask, + position_ids, + token_type_ids=prefix_token_type_ids, + ) last_logits = logits[0, -1, :] # (vocab_size,) golden = _extract_logits_golden(last_logits) @@ -304,9 +321,19 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: model_device = _get_model_device(model, device) gen_ids = torch.from_numpy(input_ids).to(model_device) max_new = case.generation_params.get("max_new_tokens", 20) + gen_kwargs: dict = {"max_new_tokens": max_new, "do_sample": False} + if prefix_token_type_ids is not None: + # HF's HrmTextForCausalLM.create_masks_for_generate applies the + # PrefixLM overlay only on the first iteration, so marking the whole + # prompt as one prefix block is exactly the model card's + # ``token_type_ids = ones_like(input_ids)`` call. + gen_kwargs["attention_mask"] = torch.from_numpy(attention_mask).to(model_device) + gen_kwargs["token_type_ids"] = torch.from_numpy(prefix_token_type_ids).to( + model_device + ) with torch.no_grad(): try: - gen_output = model.generate(gen_ids, max_new_tokens=max_new, do_sample=False) + gen_output = model.generate(gen_ids, **gen_kwargs) except ValueError as e: # All-attention GraniteMoeHybrid variants (e.g. granite-4.0-1b) # trip transformers' hybrid Mamba/attention generation cache, @@ -316,9 +343,7 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: # the (slower) cache-free path. if "has_previous_state" not in str(e): raise - gen_output = model.generate( - gen_ids, max_new_tokens=max_new, do_sample=False, use_cache=False - ) + gen_output = model.generate(gen_ids, use_cache=False, **gen_kwargs) generated_ids = gen_output[0, seq_len:].cpu().numpy() save_golden_ref( diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 5355cc873..6daf5f218 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -1814,6 +1814,19 @@ def __post_init__(self): def from_transformers(cls, config, parent_config=None) -> HrmTextConfig: base = ArchitectureConfig.from_transformers(config, parent_config) fields = _shallow_fields(base) + # ``HrmTextRotaryEmbedding`` is unconditional upstream, so this family + # always uses RoPE. A raw ``config.json`` only carries the default + # ``rope_theta`` of 10000.0, which the generic extractor deliberately + # ignores as a RoPE signal (NoPE models inherit it as dead config + # data), leaving ``rope_type=None`` and silently exporting a + # position-free graph. Pin the default RoPE explicitly instead. + if fields.get("rope_type") is None: + fields["rope_type"] = "default" + if fields.get("rope_theta") is None: + fields["rope_theta"] = 10_000.0 + if fields.get("partial_rotary_factor") is None: + # Upstream rotates the full head_dim. + fields["partial_rotary_factor"] = 1.0 return cls( **fields, H_cycles=_as_int_or_default(getattr(config, "H_cycles", None), 2), diff --git a/src/mobius/_testing/generation.py b/src/mobius/_testing/generation.py index 2cb48ffbe..271ce31b2 100644 --- a/src/mobius/_testing/generation.py +++ b/src/mobius/_testing/generation.py @@ -15,6 +15,11 @@ from mobius._configs import ArchitectureConfig from mobius._testing.ort_inference import OnnxModelSession +from mobius._testing.prefix_lm import ( + TOKEN_TYPE_IDS, + generated_token_type_ids, + prompt_token_type_ids, +) class OnnxGenerator: @@ -109,6 +114,12 @@ def _kv_dtype(name: str) -> np.dtype: all_ids = input_ids.copy() + # PrefixLM models (HRM-Text) take an extra ``token_type_ids`` input: + # the whole prompt is one bidirectional prefix block (1), every + # generated token is causal (0). Mirrors the model card's documented + # ``token_type_ids = ones_like(input_ids)`` prefill call. + uses_token_type_ids = TOKEN_TYPE_IDS in set(self.session.input_names) + # First step: process the full prompt cur_input_ids = input_ids past_seq_len = 0 @@ -128,6 +139,12 @@ def _kv_dtype(name: str) -> np.dtype: "position_ids": position_ids, **past_kv, } + if uses_token_type_ids: + feeds[TOKEN_TYPE_IDS] = ( + prompt_token_type_ids(cur_input_ids) + if past_seq_len == 0 + else generated_token_type_ids(cur_input_ids) + ) outputs = self.session.run(feeds) diff --git a/src/mobius/_testing/prefix_lm.py b/src/mobius/_testing/prefix_lm.py new file mode 100644 index 000000000..da894db27 --- /dev/null +++ b/src/mobius/_testing/prefix_lm.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""PrefixLM ``token_type_ids`` contract shared by every test/golden path. + +Models whose exported graph declares ``requires_token_type_ids`` were +pre-trained with a PrefixLM mask: prompt tokens attend bidirectionally to each +other, generated tokens attend causally. ``sapientinc/HRM-Text-1B``'s model +card spells the correct call out explicitly — mark the *entire* prompt as one +bidirectional prefix block (``token_type_ids = ones_like(input_ids)``) and +leave generated positions at ``0``. + +Both sides of every comparison must use the same rule or the numbers are +meaningless, so the rule lives here exactly once: the HuggingFace reference is +driven from :func:`model_type_uses_prefix_lm` and the exported graph from +:func:`graph_uses_prefix_lm`, and both hand out the same tensors. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + import onnx_ir as ir + +TOKEN_TYPE_IDS = "token_type_ids" + + +def model_type_uses_prefix_lm(model_type: str) -> bool: + """Whether *model_type*'s registered class consumes ``token_type_ids``.""" + from mobius._registry import registry + + try: + model_cls = registry.get(model_type) + except Exception: + return False + return bool(getattr(model_cls, "requires_token_type_ids", False)) + + +def graph_uses_prefix_lm(model: ir.Model) -> bool: + """Whether an exported graph declares a ``token_type_ids`` input.""" + return any(value.name == TOKEN_TYPE_IDS for value in model.graph.inputs) + + +def prompt_token_type_ids(input_ids: np.ndarray) -> np.ndarray: + """``token_type_ids`` for a prompt that is one bidirectional prefix block.""" + return np.ones_like(input_ids, dtype=np.int64) + + +def generated_token_type_ids(input_ids: np.ndarray) -> np.ndarray: + """``token_type_ids`` for generated positions: causal, never prefix.""" + return np.zeros_like(input_ids, dtype=np.int64) diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index e3f71783e..78332449a 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -366,6 +366,7 @@ def torch_forward( attention_mask: np.ndarray, position_ids: np.ndarray, past_key_values: object | None = None, + token_type_ids: np.ndarray | None = None, ) -> tuple[np.ndarray, object]: """Run a single forward pass on a HuggingFace causal LM model. @@ -376,6 +377,10 @@ def torch_forward( position_ids: [batch, seq_len] int64 numpy array. past_key_values: Optional list of (key, value) numpy array tuples, or an opaque HuggingFace Cache for hybrid recurrent models. + token_type_ids: Optional [batch, seq_len] int64 PrefixLM block marker + (``1`` = bidirectional prefix position). Only forwarded when the + model's ``forward`` actually accepts it, so passing it for a + non-PrefixLM model is inert rather than an error. Returns: Tuple of logits and either a list of KV numpy tuples or an opaque @@ -401,6 +406,9 @@ def torch_forward( if "position_ids" in fwd_sig.parameters: kwargs["position_ids"] = pos_t + if token_type_ids is not None and "token_type_ids" in fwd_sig.parameters: + kwargs["token_type_ids"] = torch.from_numpy(token_type_ids).to(device) + if past_key_values is not None: from transformers.cache_utils import Cache, DynamicCache diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index 0c2f3309a..e6394e2fd 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -273,6 +273,13 @@ def __init__( self._rope_interleave = config.rope_interleave # Gemma2-style logit soft-capping; 0.0 means disabled. self._softcap = getattr(config, "attn_logit_softcapping", 0.0) or 0.0 + # Whether the Attention op applies its own causal mask. Subclasses that + # feed a float additive bias already encoding causality *and* some + # bidirectional unmasking (e.g. a PrefixLM / vision-block overlay) set + # this to 0 so the built-in mask does not cancel that unmasking. It is a + # static per-graph property, so it lives on the module rather than being + # threaded through every forward signature. + self._is_causal = 1 self.q_proj = linear_class( self.hidden_size, @@ -395,6 +402,7 @@ def forward( scale=self.scaling, softcap=self._softcap, static_cache=static_cache, + is_causal=self._is_causal, ) attn_output = self._post_attention(op, attn_output, hidden_states) diff --git a/src/mobius/models/hrm_text.py b/src/mobius/models/hrm_text.py index e4c014eb5..2a3ddeb10 100644 --- a/src/mobius/models/hrm_text.py +++ b/src/mobius/models/hrm_text.py @@ -53,9 +53,10 @@ Embedding, Linear, ScaleFreeRMSNorm, + create_attention_bias, initialize_rope, ) -from mobius.models.base import CausalLMModel, TextModel +from mobius.models.base import CausalLMModel, TextModel, _retain_last_sequence_token if TYPE_CHECKING: import onnx_ir as ir @@ -88,6 +89,12 @@ def __init__(self, config: ArchitectureConfig, linear_class: type | None = None) self.num_attention_heads * self.head_dim, bias=config.attn_qkv_bias, ) + # Under PrefixLM the model feeds a float additive bias that already + # bakes in causal + padding + the bidirectional prefix overlay, so the + # Attention op must not re-apply its own causal mask (that would cancel + # the future-position unmasking). Mirrors Gemma4's vision-block overlay. + if getattr(config, "prefix_lm", False): + self._is_causal = 0 def _post_attention( self, @@ -193,6 +200,7 @@ def __init__(self, config: ArchitectureConfig): self._l_cycles = int(config.L_cycles) self._layers_per_stack = _resolve_layers_per_stack(config) self._embedding_scale = float(config.embedding_scale) + self._prefix_lm = bool(getattr(config, "prefix_lm", False)) self.embed_tokens = Embedding( config.vocab_size, config.hidden_size, config.pad_token_id @@ -204,6 +212,78 @@ def __init__(self, config: ArchitectureConfig): # broadcast against (B, S, hidden) — equivalent to HF's ``expand_as``. self.z_L_init = nn.Parameter([config.hidden_size], dtype=config.dtype) + def _build_prefix_lm_context( + self, + op: OpBuilder, + *, + input_ids: ir.Value, + attention_mask: ir.Value | None, + position_ids: ir.Value, + token_type_ids: ir.Value | None, + ) -> tuple[ir.Value, tuple]: + """Build the PrefixLM float attention bias and the RoPE embeddings. + + Reproduces upstream ``HrmTextModel.forward``, which converts + ``token_type_ids`` into ``block_sequence_ids = where(tt == 1, 0, -1)`` + and hands it to ``create_causal_mask``. HuggingFace's + ``blockwise_overlay`` unmasks a ``(q, kv)`` pair when + ``(q_group == kv_group) & (q_group >= 0)``, OR-ed onto the causal mask; + :func:`~mobius.components.create_attention_bias` implements exactly + that, so prompt tokens (``token_type_ids == 1``, block ``0``) attend to + each other bidirectionally while everything else stays causal. + + Upstream gates the overlay on ``is_first_iteration``. The exported + graph reproduces that behaviour *by data* rather than by branching: + a generated token is fed with ``token_type_ids == 0`` -> block ``-1``, + and the ``q_group >= 0`` guard makes the overlay a pure no-op for it, + leaving plain causal attention on every decode step. Passing all-zero + ``token_type_ids`` therefore reproduces HuggingFace's + ``token_type_ids=None`` (fully causal) path exactly, so one graph + serves both contracts. + + The bias bakes in causal + padding + prefix-block masking, so + :class:`HrmTextAttention` runs with ``is_causal=0``; that also rules out + ``GroupQueryAttention``, whose mask model is causal/local-window only. + """ + 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." + ) + + position_embeddings = self.rotary_emb(op, position_ids) + + # block_sequence_ids: (B, S_q) int64 — 0 for prefix tokens, -1 otherwise. + prefix_marker = op.Constant(value_ints=[1]) + block_id = op.Constant(value_ints=[0]) + non_block_id = op.Constant(value_ints=[-1]) + block_sequence_ids = op.Where( + op.Equal(token_type_ids, prefix_marker), block_id, non_block_id + ) + + # (B, 1, S_q, past + S_q) additive float bias. + attention_bias = create_attention_bias( + op, + input_ids=input_ids, + attention_mask=attention_mask, + dtype=self._dtype, + block_sequence_ids=block_sequence_ids, + ) + return attention_bias, position_embeddings + def forward( self, op: OpBuilder, @@ -212,6 +292,7 @@ def forward( position_ids: ir.Value, past_key_values: list | None = None, inputs_embeds: ir.Value | None = None, + token_type_ids: ir.Value | None = None, ): if inputs_embeds is not None: hidden_states = inputs_embeds @@ -221,14 +302,23 @@ def forward( # z_H — slow / high-level state: (B, S, hidden) z_high = op.Mul(hidden_states, self._embedding_scale) - attention_bias, position_embeddings = self._build_attention_context( - op, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - hidden_states=z_high, - past_key_values=past_key_values, - ) + if self._prefix_lm: + attention_bias, position_embeddings = self._build_prefix_lm_context( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + token_type_ids=token_type_ids, + ) + else: + attention_bias, position_embeddings = self._build_attention_context( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + hidden_states=z_high, + past_key_values=past_key_values, + ) num_invocations = self._h_cycles * (self._l_cycles + 1) expected_slots = num_invocations * self._layers_per_stack @@ -285,11 +375,18 @@ class HrmTextCausalLMModel(CausalLMModel): sigmoid-gated attention output, and scaled token embeddings. Inputs: ``input_ids``, ``attention_mask``, ``position_ids``, - ``past_key_values``. Outputs: ``logits`` and one present KV pair per - unique attention invocation of the recurrence. + ``past_key_values``, plus ``token_type_ids`` when ``config.prefix_lm`` is + set (the pre-training PrefixLM mask: ``1`` marks a prompt position that + attends bidirectionally within the prefix block, anything else stays + causal). Outputs: ``logits`` and one present KV pair per unique attention + invocation of the recurrence. """ config_class: type = HrmTextConfig + # Class-level default so test/golden harnesses can tell, without building + # the model, that this architecture's graph carries a ``token_type_ids`` + # input. ``__init__`` narrows it per instance from ``config.prefix_lm``. + requires_token_type_ids: bool = True def __init__(self, config: ArchitectureConfig): nn.Module.__init__(self) @@ -298,6 +395,32 @@ def __init__(self, config: ArchitectureConfig): self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight + # Task hook: ask CausalLMTask for the extra ``token_type_ids`` graph + # input. Only meaningful under PrefixLM; a purely causal HRM-Text + # export keeps the standard input set. + self.requires_token_type_ids = bool(getattr(config, "prefix_lm", False)) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value | None, + position_ids: ir.Value, + past_key_values: list | None = None, + token_type_ids: ir.Value | None = None, + ): + hidden_states, present_key_values = self.model( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + token_type_ids=token_type_ids, + ) + # Honour prefill-prefix pruning exactly like CausalLMModel.forward. + hidden_states = _retain_last_sequence_token(op, hidden_states) + logits = self.lm_head(op, hidden_states) + return logits, present_key_values def preprocess_weights( self, state_dict: dict[str, torch.Tensor] diff --git a/src/mobius/models/hrm_text_test.py b/src/mobius/models/hrm_text_test.py index 836eb87b3..c7576ae22 100644 --- a/src/mobius/models/hrm_text_test.py +++ b/src/mobius/models/hrm_text_test.py @@ -146,6 +146,28 @@ def test_explicit_embedding_scale_is_preserved(): assert config.embedding_scale == pytest.approx(39.191835884530846) +def test_raw_json_config_still_gets_rope(): + """A raw config.json must not silently export a position-free graph. + + ``HrmTextRotaryEmbedding`` is unconditional upstream, but the raw + ``config.json`` only carries the *default* ``rope_theta`` of 10000.0, which + the generic extractor ignores as a RoPE signal. + """ + config = _mobius_config() + assert config.rope_type == "default" + assert config.rope_theta == pytest.approx(10_000.0) + module, _ = _build_package(config) + assert module.model.rotary_emb is not None + + +def test_trusted_hf_config_rope_matches_raw_json(): + hf_config = _hf_config() + trusted = HrmTextConfig.from_transformers(hf_config) + raw = _mobius_config() + assert trusted.rope_type == raw.rope_type == "default" + assert trusted.rope_theta == pytest.approx(raw.rope_theta) + + def test_config_rejects_non_positive_cycles(): with pytest.raises(ValueError, match="positive H_cycles"): HrmTextConfig.from_transformers(_raw_json_config(L_cycles=0)) @@ -188,6 +210,53 @@ def test_graph_exposes_one_cache_slot_per_attention_invocation(): assert f"past_key_values.{_TOTAL_SLOTS}.key" not in input_names +def test_prefix_lm_graph_declares_token_type_ids(): + config = _mobius_config() + assert config.prefix_lm is True + module, pkg = _build_package(config) + assert module.requires_token_type_ids is True + input_names = [value.name for value in pkg["model"].graph.inputs] + assert "token_type_ids" in input_names + # Declared next to the other per-position inputs, before the cache. + assert input_names[:4] == [ + "input_ids", + "attention_mask", + "position_ids", + "token_type_ids", + ] + + +def test_causal_only_config_has_no_token_type_ids_input(): + """``prefix_lm=False`` keeps the standard causal input set (and GQA path).""" + config = _mobius_config(prefix_lm=False) + module, pkg = _build_package(config) + assert module.requires_token_type_ids is False + assert "token_type_ids" not in {value.name for value in pkg["model"].graph.inputs} + + +def test_prefix_lm_forward_requires_token_type_ids(): + """A direct forward() call must not silently fall back to causal masking.""" + config = _mobius_config() + module = HrmTextCausalLMModel(config) + with pytest.raises(ValueError, match="token_type_ids"): + get_task("text-generation").build(_StripTokenTypeIds(module), config) + + +class _StripTokenTypeIds: + """Wrap a module so the task's ``token_type_ids`` never reaches forward().""" + + def __init__(self, module): + self._module = module + self.requires_token_type_ids = False + + def __getattr__(self, name): + return getattr(self._module, name) + + def __call__(self, op, **kwargs): + kwargs.pop("token_type_ids", None) + return self._module(op, **kwargs) + + def test_stack_weights_are_shared_across_recurrence_steps(): config = _mobius_config() _, pkg = _build_package(config) @@ -299,11 +368,19 @@ def test_preprocess_weights_rejects_wrong_fused_width(): def test_recurrence_matches_huggingface_prefill_and_decode(): - """Prefill *and* a cached decode step must match upstream. - - The decode step is the part that pins the KV-cache slot layout: an - off-by-one stack ordering still passes a cacheless prefill comparison but - reads the wrong slots on the second step. + """PrefixLM prefill, causal fallback, and cached decode must match upstream. + + Three things are pinned here: + + * **PrefixLM prefill** — ``token_type_ids == 1`` over the whole prompt must + reproduce upstream's ``block_sequence_ids = where(tt == 1, 0, -1)`` + bidirectional overlay. + * **Causal fallback** — all-zero ``token_type_ids`` through the *same* + graph must reproduce upstream's ``token_type_ids=None`` path, and the gap + between the two modes must be the same size on both sides. Without that + second half a graph that ignored the overlay entirely would still pass. + * **Cached decode** — the KV-cache slot layout, plus the fact that a + generated token leaves the prefix block and attends causally. """ transformers = pytest.importorskip("transformers") torch.manual_seed(11) @@ -316,45 +393,72 @@ def test_recurrence_matches_huggingface_prefill_and_decode(): apply_weights(pkg["model"], module.preprocess_weights(dict(hf_model.state_dict()))) rng = np.random.default_rng(11) - input_ids = rng.integers(1, _VOCAB, size=(1, 5)).astype(np.int64) + input_ids = rng.integers(1, _VOCAB, size=(1, 6)).astype(np.int64) attention_mask = np.ones_like(input_ids) position_ids = np.arange(input_ids.shape[1], dtype=np.int64)[np.newaxis, :] session = OnnxModelSession(pkg["model"]) try: - feeds: dict[str, np.ndarray] = { - "input_ids": input_ids, - "attention_mask": attention_mask, - "position_ids": position_ids, - } - empty = np.zeros((1, _HEADS, 0, _HEAD_DIM), dtype=np.float32) - for slot in range(_TOTAL_SLOTS): - feeds[f"past_key_values.{slot}.key"] = empty - feeds[f"past_key_values.{slot}.value"] = empty - prefill = session.run(feeds) + + def _prefill(token_type_ids: np.ndarray) -> dict[str, np.ndarray]: + feeds: dict[str, np.ndarray] = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "token_type_ids": token_type_ids, + } + empty = np.zeros((1, _HEADS, 0, _HEAD_DIM), dtype=np.float32) + for slot in range(_TOTAL_SLOTS): + feeds[f"past_key_values.{slot}.key"] = empty + feeds[f"past_key_values.{slot}.value"] = empty + return session.run(feeds) + + onnx_prefix = _prefill(np.ones_like(input_ids)) + onnx_causal = _prefill(np.zeros_like(input_ids)) with torch.no_grad(): - hf_prefill = hf_model( + hf_prefix = hf_model( input_ids=torch.from_numpy(input_ids), attention_mask=torch.from_numpy(attention_mask), position_ids=torch.from_numpy(position_ids), + token_type_ids=torch.ones_like(torch.from_numpy(input_ids)), use_cache=True, ) + hf_causal_logits = hf_model( + input_ids=torch.from_numpy(input_ids), + attention_mask=torch.from_numpy(attention_mask), + position_ids=torch.from_numpy(position_ids), + use_cache=False, + ).logits.numpy() + hf_prefix_logits = hf_prefix.logits.numpy() + + np.testing.assert_allclose( + onnx_prefix["logits"], hf_prefix_logits, rtol=1e-3, atol=1e-3 + ) np.testing.assert_allclose( - prefill["logits"], hf_prefill.logits.numpy(), rtol=1e-3, atol=1e-3 + onnx_causal["logits"], hf_causal_logits, rtol=1e-3, atol=1e-3 ) - next_id = np.array([[int(hf_prefill.logits[0, -1].argmax())]], dtype=np.int64) + hf_delta = float(np.abs(hf_prefix_logits - hf_causal_logits).max()) + onnx_delta = float(np.abs(onnx_prefix["logits"] - onnx_causal["logits"]).max()) + assert hf_delta > 1e-4, f"HF PrefixLM overlay had no effect ({hf_delta})" + assert onnx_delta == pytest.approx(hf_delta, rel=1e-2, abs=1e-5) + + # Cached decode: the generated token is NOT part of the prefix block. + next_id = np.array([[int(hf_prefix_logits[0, -1].argmax())]], dtype=np.int64) decode_mask = np.ones((1, input_ids.shape[1] + 1), dtype=np.int64) decode_pos = np.array([[input_ids.shape[1]]], dtype=np.int64) decode_feeds: dict[str, np.ndarray] = { "input_ids": next_id, "attention_mask": decode_mask, "position_ids": decode_pos, + "token_type_ids": np.zeros_like(next_id), } for slot in range(_TOTAL_SLOTS): - decode_feeds[f"past_key_values.{slot}.key"] = prefill[f"present.{slot}.key"] - decode_feeds[f"past_key_values.{slot}.value"] = prefill[f"present.{slot}.value"] + decode_feeds[f"past_key_values.{slot}.key"] = onnx_prefix[f"present.{slot}.key"] + decode_feeds[f"past_key_values.{slot}.value"] = onnx_prefix[ + f"present.{slot}.value" + ] decode = session.run(decode_feeds) with torch.no_grad(): @@ -362,7 +466,7 @@ def test_recurrence_matches_huggingface_prefill_and_decode(): input_ids=torch.from_numpy(next_id), attention_mask=torch.from_numpy(decode_mask), position_ids=torch.from_numpy(decode_pos), - past_key_values=hf_prefill.past_key_values, + past_key_values=hf_prefix.past_key_values, use_cache=True, ) np.testing.assert_allclose( diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 41758389e..6d4001c75 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -43,6 +43,9 @@ class CausalLMTask(ModelTask): - input_ids: [batch, sequence_len] INT64 - attention_mask: [batch, total_seq_len] INT64 - position_ids: [batch, sequence_len] INT64 + - token_type_ids: [batch, sequence_len] INT64 — only for modules + that set ``requires_token_type_ids = True`` (PrefixLM models + such as HRM-Text); absent for every other model. - past_key_values.{i}.key: [batch, num_kv_heads, past_seq_len, head_dim] - past_key_values.{i}.value: [batch, num_kv_heads, past_seq_len, head_dim] Outputs: @@ -133,6 +136,10 @@ def build( # --- Inputs common to both modes --- input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + # Optional module-declared extra graph inputs, forwarded as forward() + # keywords only for the modules that ask for them. + extra_module_kwargs: dict[str, ir.Value] = {} + # --- Cache setup (static vs dynamic) --- if static: attention_mask = None @@ -165,6 +172,16 @@ def build( "position_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] ) + # PrefixLM-style models declare that they consume a per-position + # ``token_type_ids`` tensor to build a bidirectional prefix block + # (HRM-Text). This is opt-in per module — exactly like the + # ``kv_cache_layer_count`` / ``static_kv_cache_specs`` hooks below — + # so no other model gains an input or a forward keyword. + if getattr(module, "requires_token_type_ids", False): + extra_module_kwargs["token_type_ids"] = builder.input( + "token_type_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) + # MLA attention: K/V heads equal q heads (no GQA reduction in # latent space). The ONNX Attention op is called with # kv_num_heads=num_attention_heads, so the KV cache must use @@ -205,6 +222,7 @@ def build( attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, + **extra_module_kwargs, ) intermediate_hidden_states: list | None = None if len(result) == 3: diff --git a/testdata/cases/causal-lm/hrm-text-1b.yaml b/testdata/cases/causal-lm/hrm-text-1b.yaml index 69e11e0cc..ef3d681d7 100644 --- a/testdata/cases/causal-lm/hrm-text-1b.yaml +++ b/testdata/cases/causal-lm/hrm-text-1b.yaml @@ -6,13 +6,14 @@ dtype: "float32" inputs: prompts: - - "Here is my poem:" + - "<|im_start|><|quad_end|><|object_ref_end|>Explain why the sky is blue.<|im_end|>" level: "L4+L5" generation: max_new_tokens: 24 do_sample: false + exact_match: true notes: >- HRM-Text 1B hierarchical recurrent decoder. Two independently-weighted @@ -20,4 +21,8 @@ notes: >- = 8 times per forward, so the graph carries 128 KV-cache slots for 16 per-stack blocks. Parameterless RMSNorm, sigmoid-gated attention output, and token embeddings scaled by 1 / initializer_range. Base model with no chat - template, so the prompt is fed through the plain tokenizer. + template, so the model card's documented prompt (including its literal + control tokens) is fed through the plain tokenizer. Both the reference and + the exported graph run the PrefixLM contract from the model card -- the whole + prompt is one bidirectional block (token_type_ids = 1), generated tokens are + causal (0); see mobius._testing.prefix_lm. diff --git a/testdata/golden/causal-lm/hrm-text-1b.json b/testdata/golden/causal-lm/hrm-text-1b.json new file mode 100644 index 000000000..672d5291e --- /dev/null +++ b/testdata/golden/causal-lm/hrm-text-1b.json @@ -0,0 +1,50 @@ +{ + "top1_id": 341, + "top2_id": 329, + "top10_ids": [ + 341, + 329, + 1783, + 2718, + 47, + 81, + 49701, + 63, + 2255, + 64 + ], + "top10_logits": [ + "0x1.08c4220000000p+4", + "0x1.98941c0000000p+3", + "0x1.5c3f860000000p+3", + "0x1.4a5e700000000p+3", + "0x1.3abd740000000p+3", + "0x1.3085480000000p+3", + "0x1.1e36a00000000p+3", + "0x1.1bb8d80000000p+3", + "0x1.166b180000000p+3", + "0x1.1353520000000p+3" + ], + "logits_summary": [ + "0x1.08c4220000000p+4", + "-0x1.96a4140000000p+2", + "-0x1.e62a31e80ec00p-1", + "0x1.08d99152a0633p+0" + ], + "input_ids": [ + 6, + 13, + 9, + 30683, + 1685, + 236, + 114, + 102, + 99, + 7884, + 322, + 3288, + 44, + 7 + ] +} diff --git a/testdata/golden/causal-lm/hrm-text-1b_generation.json b/testdata/golden/causal-lm/hrm-text-1b_generation.json new file mode 100644 index 000000000..a87ff17fa --- /dev/null +++ b/testdata/golden/causal-lm/hrm-text-1b_generation.json @@ -0,0 +1,31 @@ +{ + "model_id": "sapientinc/HRM-Text-1B", + "prompt": "<|im_start|><|quad_end|><|object_ref_end|>Explain why the sky is blue.<|im_end|>", + "generated_tokens": [ + 341, + 7884, + 2166, + 3288, + 1794, + 309, + 236, + 114, + 102, + 99, + 19173, + 301, + 10404, + 431, + 1817, + 6264, + 44, + 1172, + 5613, + 15436, + 301, + 2282, + 395, + 23618 + ], + "generated_text": "The sky appears blue due to the scattering of sunlight by air molecules. Shorter wavelengths of light (blue" +} diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index 4da2e5c6e..17b6e01f0 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -56,6 +56,7 @@ ) from mobius._testing.ort_inference import OnnxModelSession from mobius._testing.parity import ParityResult, compare_golden +from mobius._testing.prefix_lm import prompt_token_type_ids @functools.cache @@ -682,9 +683,20 @@ def _prepare_prefill_feeds( "position_ids": np.arange(seq_len, dtype=np.int64).reshape(1, -1), } - # Provide token_type_ids for models that need it (BERT, ALBERT, etc.) + # Provide token_type_ids for models that need it. + # + # Two different meanings share this input name: + # * BERT/ALBERT-style segment ids — a single-segment prompt is all zeros. + # * PrefixLM block markers (HRM-Text, ``config.prefix_lm``) — the whole + # prompt is one bidirectional prefix block, i.e. all ones. This must + # match the contract the golden reference was generated with (see + # ``mobius._testing.prefix_lm``), otherwise the graph runs a fully + # causal mask against a PrefixLM golden. if "token_type_ids" in session.input_names: - feeds["token_type_ids"] = np.zeros_like(input_ids) + if getattr(config, "prefix_lm", False): + feeds["token_type_ids"] = prompt_token_type_ids(input_ids) + else: + feeds["token_type_ids"] = np.zeros_like(input_ids) # Fill KV cache inputs with zero-length tensors. # Shape: (batch=1, num_kv_heads, past_seq_len=0, head_dim) diff --git a/tests/integration_test.py b/tests/integration_test.py index 233b97df0..e5e4f36c9 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -40,6 +40,10 @@ ) from mobius._testing.generation import OnnxGenerator, torch_generate_greedy from mobius._testing.ort_inference import OnnxModelSession +from mobius._testing.prefix_lm import ( + generated_token_type_ids, + prompt_token_type_ids, +) from mobius._testing.torch_reference import ( load_torch_model, load_torch_multimodal_model, @@ -246,21 +250,30 @@ def test_minicpmv4_6_real_weight_vision_parity(): _HRM_TEXT_MODEL_ID = "sapientinc/HRM-Text-1B" _HRM_TEXT_REVISION = "9f082d68b8cd0ebc56e33f1c88c45609174c272c" +# The model card's documented prompt. HRM-Text-1B is a base model with no chat +# template, so its control tokens are typed literally through the plain +# tokenizer rather than via ``apply_chat_template``. +_HRM_TEXT_PROMPT = ( + "<|im_start|><|quad_end|><|object_ref_end|>Explain why the sky is blue.<|im_end|>" +) @pytest.mark.integration @pytest.mark.integration_slow def test_hrm_text_1b_real_weight_prefill_and_decode_parity(): - """Real HRM-Text-1B weights: prefill and cached decode must match HF. + """Real HRM-Text-1B weights: PrefixLM prefill and cached decode match HF. - The cached decode step is the load-bearing half: HRM-Text drives 128 + Runs the model card's documented contract on both sides — the whole prompt + marked as one bidirectional prefix block via ``token_type_ids``. + + The cached decode step is load-bearing twice over. HRM-Text drives 128 KV-cache slots from only 32 blocks of weights, so a slot-ordering mistake - in the H/L recurrence still passes a prefill-only comparison but reads the - wrong cache entries on the second step. + in the H/L recurrence still passes a prefill-only comparison; and the + generated token must leave the prefix block, so a graph that kept the + bidirectional overlay switched on would diverge here too. - ``sapientinc/HRM-Text-1B`` is a base model with no chat template, so the - prompt goes through the plain tokenizer. The pinned revision is forwarded - to the config, tokenizer, weights, and the mobius build alike. + The pinned revision is forwarded to the config, tokenizer, weights, and the + mobius build alike. """ import gc @@ -276,13 +289,15 @@ def test_hrm_text_1b_real_weight_prefill_and_decode_parity(): assert config.num_hidden_layers == ( config.num_layers_per_stack * config.H_cycles * (config.L_cycles + 1) ) + assert config.prefix_lm is True tokenizer = AutoTokenizer.from_pretrained(_HRM_TEXT_MODEL_ID, revision=_HRM_TEXT_REVISION) - tokens = tokenizer("Here is my poem:", return_tensors="np") + tokens = tokenizer(_HRM_TEXT_PROMPT, return_tensors="np") input_ids = tokens["input_ids"].astype(np.int64) attention_mask = tokens["attention_mask"].astype(np.int64) seq_len = input_ids.shape[1] position_ids = np.arange(seq_len, dtype=np.int64)[np.newaxis, :] + token_type_ids = prompt_token_type_ids(input_ids) torch_model = AutoModelForCausalLM.from_pretrained( _HRM_TEXT_MODEL_ID, @@ -294,9 +309,27 @@ def test_hrm_text_1b_real_weight_prefill_and_decode_parity(): input_ids=torch.from_numpy(input_ids), attention_mask=torch.from_numpy(attention_mask), position_ids=torch.from_numpy(position_ids), + token_type_ids=torch.from_numpy(token_type_ids), use_cache=True, ) + # Same prompt with the overlay disabled, to prove the exported graph + # really applies PrefixLM instead of quietly running plain causal. + hf_causal_logits = ( + torch_model( + input_ids=torch.from_numpy(input_ids), + attention_mask=torch.from_numpy(attention_mask), + position_ids=torch.from_numpy(position_ids), + use_cache=False, + ) + .logits.float() + .numpy() + ) hf_prefill_logits = hf_prefill.logits.float().numpy() + prefix_vs_causal = float(np.abs(hf_prefill_logits - hf_causal_logits).max()) + assert prefix_vs_causal > 1.0, ( + "PrefixLM and causal masking must differ substantially for this " + f"checkpoint, got max abs diff {prefix_vs_causal}" + ) next_token = np.array([[int(hf_prefill_logits[0, -1].argmax())]], dtype=np.int64) decode_mask = np.ones((1, seq_len + 1), dtype=np.int64) @@ -328,6 +361,7 @@ def test_hrm_text_1b_real_weight_prefill_and_decode_parity(): name for name in session.input_names if name.startswith("past_key_values.") ] assert len(cache_inputs) == 2 * config.num_hidden_layers + assert "token_type_ids" in session.input_names feeds = _make_prefill_feeds(config, input_ids, attention_mask, position_ids) onnx_prefill = session.run(feeds) @@ -338,6 +372,13 @@ def test_hrm_text_1b_real_weight_prefill_and_decode_parity(): ) onnx_decode = session.run(decode_feeds) assert_logits_close(onnx_decode["logits"], hf_decode_logits, rtol=1e-3, atol=1e-3) + + # All-zero token_type_ids must reproduce HF's causal + # (``token_type_ids=None``) path through the very same graph. + causal_feeds = dict(feeds) + causal_feeds["token_type_ids"] = generated_token_type_ids(input_ids) + onnx_causal = session.run(causal_feeds) + assert_logits_close(onnx_causal["logits"], hf_causal_logits, rtol=1e-3, atol=1e-3) finally: session.close() @@ -468,6 +509,9 @@ def _make_prefill_feeds(config, input_ids, attention_mask, position_ids): "attention_mask": attention_mask, "position_ids": position_ids, } + # PrefixLM exports (HRM-Text) take the prompt as one bidirectional block. + if getattr(config, "prefix_lm", False): + feeds["token_type_ids"] = prompt_token_type_ids(input_ids) for i in range(config.num_hidden_layers): layer_types = config.layer_types or [] layer_type = layer_types[i] if i < len(layer_types) else "full_attention" @@ -497,6 +541,9 @@ def _make_decode_feeds( "attention_mask": decode_attention_mask, "position_ids": decode_position_ids, } + # A generated token is never part of the bidirectional prefix block. + if getattr(config, "prefix_lm", False): + feeds["token_type_ids"] = generated_token_type_ids(decode_input_ids) for i in range(config.num_hidden_layers): layer_types = config.layer_types or [] layer_type = layer_types[i] if i < len(layer_types) else "full_attention" diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index ef874c158..ec8e6fefc 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -1228,6 +1228,12 @@ def test_synthetic_parity(model_type: str, config_overrides: dict): "attention_mask": attention_mask, "position_ids": position_ids, } + # PrefixLM models (HRM-Text) take an extra ``token_type_ids`` input. The HF + # reference above runs without ``token_type_ids``, i.e. fully causal, so + # feed all zeros here — block id -1 makes the bidirectional overlay a no-op + # and the two sides compare the same masking contract. + if "token_type_ids" in {inp.name for inp in onnx_model.graph.inputs}: + feeds["token_type_ids"] = np.zeros_like(input_ids) # Add zero-valued past KV cache feeds with correct shapes: # batch=1, past_sequence_len=0, other dims from model spec for inp in onnx_model.graph.inputs: