From bf4380c4168f0a7594c6cdcacc44a62a34c30d65 Mon Sep 17 00:00:00 2001 From: Synapticode Agent Date: Sat, 13 Jun 2026 12:50:44 +1000 Subject: [PATCH] feat(adapters): HRM-Text dual-timescale recurrent adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First recurrent architecture in the adapter registry. HRM-Text (HrmTextForCausalLM, Sapient HRM-Text) is dual-timescale: two separately- parameterised H/L block stacks (slow z_H / fast z_L) that share a block design but own distinct weights; per forward the L-stack executes 6x and the H-stack 2x. The safetensors stores fused attn.gqkv_proj (gated QKV) and mlp.gate_up_proj — the converter reads these directly (128 fused ternary tensors for the 1B, 64 H + 64 L; not the 256 split state_dict). MagicNorm is parameterless, so the only protected tensors are the embeddings, the untied LM head, and the 1-D recurrent init z_L_init (named explicitly). classify_weight reuses the standard substring+ dimensionality policy; a stack_of() helper tags H/L/shared for per-stack reporting of the asymmetric reuse leverage. Pure addition: one new adapter file + one registry entry (registry canonical-set test updated to match). No existing adapter or converter logic touched. Verified against sapientinc/HRM-Text-1B: dry-run + full convert both 128 ternary / 3 fp16 / 83.0% ternary / 3.07x vs bf16 / ~63s, matching the runtime-registered prototype byte-for-byte (probe HRM_PROBE_20260613T011321Z). Footprint only — no throughput (recurrence-forward runtime integration scoped, not run) and no coherence claim (Lane B held under program QAT hold). Co-Authored-By: Claude Opus 4.8 --- src/terncore/adapters/__init__.py | 2 +- src/terncore/adapters/hrm_text.py | 168 ++++++++++++++++++++++++++++++ tests/test_adapters_registry.py | 6 +- tests/test_hrm_text_adapter.py | 151 +++++++++++++++++++++++++++ 4 files changed, 324 insertions(+), 3 deletions(-) create mode 100644 src/terncore/adapters/hrm_text.py create mode 100644 tests/test_hrm_text_adapter.py diff --git a/src/terncore/adapters/__init__.py b/src/terncore/adapters/__init__.py index 0804659..a7e992c 100644 --- a/src/terncore/adapters/__init__.py +++ b/src/terncore/adapters/__init__.py @@ -18,7 +18,7 @@ _REGISTRY: dict[str, type["ArchitectureAdapter"]] = {} -_KNOWN_ADAPTERS = ["llama", "gemma3", "gemma4", "gemma4_unified", "phi3", "qwen3", "qwen3_moe", "kokoro"] +_KNOWN_ADAPTERS = ["llama", "gemma3", "gemma4", "gemma4_unified", "phi3", "qwen3", "qwen3_moe", "kokoro", "hrm_text"] def register(name: str): diff --git a/src/terncore/adapters/hrm_text.py b/src/terncore/adapters/hrm_text.py new file mode 100644 index 0000000..3dc91e7 --- /dev/null +++ b/src/terncore/adapters/hrm_text.py @@ -0,0 +1,168 @@ +""" +HRM-Text dual-timescale recurrent architecture adapter. + +Maps Sapient Intelligence's HRM-Text (``HrmTextForCausalLM``, +``model_type="hrm_text"``) HuggingFace weight names to tern-core's +internal conversion schema. HRM-Text is the first **recurrent** +architecture tern-core ingests — it is not a stock transformer. + +Architecture (probe 2026-06-13, ``sapientinc/HRM-Text-1B``): +- **Dual-timescale recurrence.** Two independently-parameterised stacks + share the same block design but own separate weights: + ``model.H_module.layers.N.*`` (slow / strategic, ``z_H``) and + ``model.L_module.layers.N.*`` (fast / execution, ``z_L``). Per forward + pass the **L-stack executes 6×** (2 H-cycles × 3 L-steps) and the + **H-stack 2×** over the same recurrent state — so the two stacks carry + identical footprint but different per-forward reuse leverage. The + conversion is symmetric; :meth:`stack_of` tags each weight H/L for + per-stack reporting (a reporting aid, not a correctness requirement — + classification is stack-agnostic). +- **Fused projections in the safetensors.** The on-disk checkpoint stores + attention as a single ``attn.gqkv_proj`` ``[4*hidden, hidden]`` (gated + Q/K/V) and the MLP as a single ``mlp.gate_up_proj`` ``[2*intermediate, + hidden]``; the HF modeling code splits these into ``q/k/v/gate`` and + ``gate/up`` on load. **tern-core reads the safetensors directly**, so it + ternarises the *fused* tensors — the conversion sees 128 ternary tensors + for HRM-Text-1B (64 H + 64 L), not the 256 of the split state_dict. Both + views cover the same parameters; the fused layout is the ground truth + for conversion and the source of the layer count in reports. +- **Parameterless MagicNorm.** Block norms carry no weights, so the only + protected tensors are the embeddings, the (untied) LM head, and the 1-D + recurrent initialiser ``model.z_L_init`` (named explicitly below for + clarity; the generic 1-D rule would also retain it). + +This is a dense, text-only adapter — no MoE expert stacking +(:meth:`expand_stacked` stays at the base ``None``) and no multimodal +components. Weight classification reuses the standard substring + +dimensionality policy: embeddings / LM head / norms / ``z_L_init`` → +FP16-retain; remaining 2-D block weights → ternary-eligible. + +Copyright (c) 2025–2026 Gamma Seeds Pte Ltd. All rights reserved. +""" + +from __future__ import annotations + +import re +from typing import Literal, Optional + +from terncore.adapters import register +from terncore.adapters.base import ( + AdapterInfo, + ArchitectureAdapter, + WeightClassification, +) + +# Matches the per-block index in both stacks +# (``model.H_module.layers.7.attn.gqkv_proj.weight`` -> 7). The H/L stack +# is *not* captured here — both stacks share ``.layers.N.`` numbering; use +# :meth:`stack_of` to disambiguate the stack. +_BLOCK_PATTERN = re.compile(r"\.layers\.(\d+)\.") + +# Fused HRM projection names, ordered by ternary tolerance (most-tolerant +# first), for reporting / priority. ``gqkv_proj`` = fused gated Q/K/V, +# ``gate_up_proj`` = fused MLP gate+up. +_PROJ_PRIORITY = [ + "gqkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", +] + +# Protection patterns. ``z_l_init`` is named explicitly (the 1-D recurrent +# state initialiser) even though the 1-D rule would retain it anyway; the +# remaining patterns cover embeddings, the untied LM head, and any norm +# weights a future HRM variant might carry. +_ALWAYS_PROTECTED = ( + "embed_tokens", + "lm_head", + "z_l_init", + "norm", + "layernorm", + "layer_norm", + "rmsnorm", + "classifier", +) + + +@register("hrm_text") +class HrmTextAdapter(ArchitectureAdapter): + """Architecture adapter for the recurrent ``HrmTextForCausalLM`` family. + + Weight classification: + 1. Embeddings, LM head, ``z_L_init``, norms → FP16-retain. + 2. 1-D weights (scalars / norms / recurrent init) → FP16-retain. + 3. All 2-D weights in the H/L block stacks (fused ``gqkv_proj`` / + ``gate_up_proj``, ``o_proj``, ``down_proj``) → ternary-eligible. + """ + + def info(self) -> AdapterInfo: + return AdapterInfo( + name="hrm_text", + architectures=["HrmTextForCausalLM"], + model_type="hrm_text", + description=( + "HRM-Text dual-timescale recurrent adapter — " + "HrmTextForCausalLM (Sapient HRM-Text). Two separately-" + "parameterised H/L block stacks (slow z_H / fast z_L), " + "fused gqkv_proj + gate_up_proj in the safetensors, " + "parameterless MagicNorm. Dense, text-only." + ), + block_pattern=_BLOCK_PATTERN, + projection_priority=list(_PROJ_PRIORITY), + protection_patterns=list(_ALWAYS_PROTECTED), + multimodal=False, + ) + + def normalize_name(self, name: str) -> str: + return name + + def stack_of(self, name: str) -> Literal["H", "L", "shared"]: + """Return the recurrence stack a weight belongs to. + + ``"H"`` for the slow/strategic stack (``H_module``), ``"L"`` for + the fast/execution stack (``L_module``), ``"shared"`` for tensors + outside both stacks (embeddings, LM head, ``z_L_init``). This is a + reporting aid — it lets ``--verbose`` / conversion reports tag each + layer by stack so the asymmetric per-forward reuse (L 6× vs H 2×) + is visible. Classification itself does not depend on the stack. + """ + if "H_module" in name: + return "H" + if "L_module" in name: + return "L" + return "shared" + + def classify_weight( + self, + name: str, + shape: Optional[list[int]] = None, + ) -> WeightClassification: + canonical = self.normalize_name(name) + name_lower = canonical.lower() + + for pattern in _ALWAYS_PROTECTED: + if pattern in name_lower: + return WeightClassification( + name=name, + canonical_name=canonical, + category="fp16_retain", + reason=f"Protected pattern: '{pattern}'", + component="language", + ) + + if shape is not None and len(shape) < 2: + return WeightClassification( + name=name, + canonical_name=canonical, + category="fp16_retain", + reason="1-D tensor (norm, scalar, or recurrent init)", + component="language", + ) + + return WeightClassification( + name=name, + canonical_name=canonical, + category="ternary_eligible", + reason="2-D weight in H/L block stack", + component="language", + ) diff --git a/tests/test_adapters_registry.py b/tests/test_adapters_registry.py index 6fe5f51..73dfee6 100644 --- a/tests/test_adapters_registry.py +++ b/tests/test_adapters_registry.py @@ -34,7 +34,7 @@ def test_get_adapter_raises_on_unknown_name_case_insensitive(): @pytest.mark.parametrize( "name", - ["llama", "gemma3", "gemma4", "gemma4_unified", "phi3", "qwen3", "qwen3_moe", "kokoro"], + ["llama", "gemma3", "gemma4", "gemma4_unified", "phi3", "qwen3", "qwen3_moe", "kokoro", "hrm_text"], ) def test_get_adapter_returns_instance_for_each_known_name(name): adapter = get_adapter(name) @@ -56,7 +56,9 @@ def test_known_adapters_is_canonical_source(): # integration of integration³ Provider³ Protocol (Phase 0 brief # ``2026-05-19_kokoro_82m_integration3_attachment_phase0.md``; # OQ-1 Option A sibling-cohort placement). + # HRM-Text adapter added 2026-06-13 — first recurrent architecture + # (dual-timescale H/L stacks); probe HRM_PROBE_20260613T011321Z. assert set(_KNOWN_ADAPTERS) == { "llama", "gemma3", "gemma4", "gemma4_unified", - "phi3", "qwen3", "qwen3_moe", "kokoro", + "phi3", "qwen3", "qwen3_moe", "kokoro", "hrm_text", } diff --git a/tests/test_hrm_text_adapter.py b/tests/test_hrm_text_adapter.py new file mode 100644 index 0000000..b8741fa --- /dev/null +++ b/tests/test_hrm_text_adapter.py @@ -0,0 +1,151 @@ +""" +Tests for the HRM-Text dual-timescale recurrent adapter +(``terncore.adapters.hrm_text``). + +Probe HRM_PROBE_20260613T011321Z (2026-06-13) confirmed +``sapientinc/HRM-Text-1B`` is ``HrmTextForCausalLM`` — two separately- +parameterised H/L block stacks with **fused** ``attn.gqkv_proj`` (gated +Q/K/V) and ``mlp.gate_up_proj`` projections in the safetensors, and +parameterless MagicNorm (no block-norm weights). The only protected +tensors are the embeddings, the untied LM head, and the 1-D recurrent +initialiser ``model.z_L_init``. This suite pins the weight-classification +policy against the real fused tensor names + shapes, the H/L stack +tagging, and the architecture routing boundary. + +Copyright (c) 2025–2026 Gamma Seeds Pte Ltd. All rights reserved. +""" + +from __future__ import annotations + +import pytest + +from terncore.adapters import get_adapter +from terncore.adapters.base import ArchitectureAdapter, ArchitectureMismatch +from terncore.adapters.hrm_text import HrmTextAdapter + +# Real fused HRM-Text-1B projection tensors (raw safetensors layout), +# one representative layer per stack. gqkv_proj = fused gated Q/K/V +# [4*hidden, hidden]; gate_up_proj = fused MLP gate+up [2*inter, hidden]. +_PROJECTIONS = { + "model.H_module.layers.0.attn.gqkv_proj.weight": [6144, 1536], + "model.H_module.layers.0.attn.o_proj.weight": [1536, 1536], + "model.H_module.layers.0.mlp.gate_up_proj.weight": [8192, 1536], + "model.H_module.layers.0.mlp.down_proj.weight": [1536, 4096], + "model.L_module.layers.0.attn.gqkv_proj.weight": [6144, 1536], + "model.L_module.layers.0.attn.o_proj.weight": [1536, 1536], + "model.L_module.layers.0.mlp.gate_up_proj.weight": [8192, 1536], + "model.L_module.layers.0.mlp.down_proj.weight": [1536, 4096], +} + +# Protected: untied embeddings + LM head, and the 1-D recurrent init. +_PROTECTED = { + "model.embed_tokens.weight": [65536, 1536], + "lm_head.weight": [65536, 1536], + "model.z_L_init": [1536], +} + + +@pytest.fixture +def adapter() -> HrmTextAdapter: + return HrmTextAdapter() + + +# ── Registry / identity ─────────────────────────────────────────────── +def test_get_adapter_returns_hrm_text(): + a = get_adapter("hrm_text") + assert isinstance(a, ArchitectureAdapter) + assert isinstance(a, HrmTextAdapter) + assert a.info().name == "hrm_text" + + +def test_info_declares_recurrent_architecture(): + info = HrmTextAdapter().info() + assert info.architectures == ["HrmTextForCausalLM"] + assert info.model_type == "hrm_text" + assert info.multimodal is False + # Dense recurrent, not MoE — no expert pattern / stacking. + assert info.expert_pattern is None + + +# ── Architecture routing boundary ───────────────────────────────────── +def test_validate_accepts_hrm_text(adapter): + adapter.validate_architecture("HrmTextForCausalLM") # must not raise + + +def test_hrm_adapter_rejects_other_architectures(adapter): + for arch in ("Qwen3ForCausalLM", "LlamaForCausalLM", "Gemma4ForConditionalGeneration"): + with pytest.raises(ArchitectureMismatch): + adapter.validate_architecture(arch) + + +def test_other_adapters_reject_hrm_text(): + """Symmetric guard: a standard-transformer adapter must not absorb HRM.""" + for name in ("qwen3", "llama", "gemma4_unified"): + with pytest.raises(ArchitectureMismatch): + get_adapter(name).validate_architecture("HrmTextForCausalLM") + + +# ── Weight classification (fused projections) ───────────────────────── +@pytest.mark.parametrize("name, shape", list(_PROJECTIONS.items())) +def test_fused_projection_weights_are_ternary_eligible(adapter, name, shape): + cls = adapter.classify_weight(name, shape) + assert cls.category == "ternary_eligible" + assert cls.component == "language" + + +@pytest.mark.parametrize("name, shape", list(_PROTECTED.items())) +def test_protected_weights_are_fp16_retained(adapter, name, shape): + cls = adapter.classify_weight(name, shape) + assert cls.category == "fp16_retain" + + +def test_z_l_init_protected_by_explicit_name(adapter): + """The 1-D recurrent initialiser is named explicitly in the protection + patterns (and the 1-D rule would retain it regardless).""" + cls = adapter.classify_weight("model.z_L_init", [1536]) + assert cls.category == "fp16_retain" + assert "z_l_init" in cls.reason.lower() + + +def test_one_dimensional_weight_retained_even_without_protected_name(adapter): + cls = adapter.classify_weight("model.H_module.layers.0.some_scale", [1536]) + assert cls.category == "fp16_retain" + assert "1-D" in cls.reason + + +# ── H/L stack tagging ───────────────────────────────────────────────── +def test_stack_of_tags_h_l_and_shared(adapter): + assert adapter.stack_of("model.H_module.layers.3.attn.gqkv_proj.weight") == "H" + assert adapter.stack_of("model.L_module.layers.3.mlp.down_proj.weight") == "L" + assert adapter.stack_of("model.embed_tokens.weight") == "shared" + assert adapter.stack_of("lm_head.weight") == "shared" + assert adapter.stack_of("model.z_L_init") == "shared" + + +def test_h_and_l_layer_zero_are_distinct_entries(adapter): + """H and L share `.layers.N.` numbering but are separate stacks — the + converter must not collide them (probe verified 64 H + 64 L).""" + h_name = "model.H_module.layers.0.attn.gqkv_proj.weight" + l_name = "model.L_module.layers.0.attn.gqkv_proj.weight" + assert h_name != l_name + assert adapter.stack_of(h_name) == "H" + assert adapter.stack_of(l_name) == "L" + assert adapter.classify_weight(h_name, [6144, 1536]).category == "ternary_eligible" + assert adapter.classify_weight(l_name, [6144, 1536]).category == "ternary_eligible" + + +# ── Integration over a full representative weight set ────────────────── +def test_get_ternary_eligible_selects_only_fused_projections(adapter): + all_shapes = {**_PROJECTIONS, **_PROTECTED} + eligible = set(adapter.get_ternary_eligible(all_shapes)) + assert eligible == set(_PROJECTIONS) + assert "model.embed_tokens.weight" not in eligible + assert "model.z_L_init" not in eligible + + +# ── Block helpers ───────────────────────────────────────────────────── +def test_block_index_and_membership(adapter): + name = "model.L_module.layers.7.attn.o_proj.weight" + assert adapter.is_block_weight(name) is True + assert adapter.block_index(name) == 7 + assert adapter.is_block_weight("model.embed_tokens.weight") is False