"""MLX implementation of HRM-Text (Sapient Intelligence Hierarchical Reasoning Model). Faithful port of `sapientinc/HRM-Text-1B` with: - 128-slot recurrent KV cache (one per cycle-step × layer) - mx.fast.scaled_dot_product_attention (Metal-optimized) - mx.fast.rope with cache-aware position offset - mx.fast.rms_norm (parameterless, weight=None) - Greedy + temperature/top-p sampling - Streaming generation - Prompt formatting with HRM condition modes Architecture (per the paper / HF model card): z_H = embed(input_ids) * embedding_scale # slow / strategic z_L = zeros # fast / execution for h in range(H_cycles): for l in range(L_cycles): z_L = L_module(z_L + z_H) # shared transformer stack z_H = H_module(z_H + z_L) # shared transformer stack logits = lm_head(z_H) Each stack invocation maps to its own KV cache slot: slot(h, l, layer) = (h * (L_cycles + 1) + l) * num_layers_per_stack + layer # L-module call (h, l) slot(h, layer) = (h * (L_cycles + 1) + L_cycles) * num_layers_per_stack + layer # H-module call (h) Weight loading reads the fused HF safetensors layout 1:1: attn.gqkv_proj [4H, H] — chunked on dim 0 as [gate, q, k, v] mlp.gate_up_proj [2I, H] — chunked on dim 0 as [gate, up] """ from __future__ import annotations import json import math from dataclasses import dataclass from pathlib import Path from typing import Iterator, Optional, Union import mlx.core as mx import mlx.core.fast as mxf import mlx.nn as nn from safetensors import safe_open # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @dataclass class HrmTextConfig: vocab_size: int hidden_size: int intermediate_size: int num_layers_per_stack: int num_attention_heads: int head_dim: int max_position_embeddings: int rms_norm_eps: float rope_theta: float H_cycles: int L_cycles: int embedding_scale: float prefix_lm: bool = True tie_word_embeddings: bool = False pad_token_id: Optional[int] = None bos_token_id: Optional[int] = None eos_token_id: Optional[int] = None @classmethod def from_dict(cls, d: dict) -> "HrmTextConfig": return cls( vocab_size=d["vocab_size"], hidden_size=d["hidden_size"], intermediate_size=d["intermediate_size"], num_layers_per_stack=d["num_hidden_layers"], # checkpoint stores per-stack count num_attention_heads=d["num_attention_heads"], head_dim=d["head_dim"], max_position_embeddings=d["max_position_embeddings"], rms_norm_eps=d["rms_norm_eps"], rope_theta=d["rope_theta"], H_cycles=d["H_cycles"], L_cycles=d["L_cycles"], embedding_scale=d["embedding_scale"], prefix_lm=d.get("prefix_lm", True), tie_word_embeddings=d.get("tie_word_embeddings", False), pad_token_id=d.get("pad_token_id"), bos_token_id=d.get("bos_token_id"), eos_token_id=d.get("eos_token_id"), ) @property def num_cache_slots(self) -> int: return self.H_cycles * (self.L_cycles + 1) * self.num_layers_per_stack # --------------------------------------------------------------------------- # KV cache # --------------------------------------------------------------------------- class _KVSlot: """Single layer/cycle KV slot with chunked-grow allocation (mlx-lm style).""" __slots__ = ("keys", "values", "offset", "step") def __init__(self, step: int = 256): self.keys: Optional[mx.array] = None self.values: Optional[mx.array] = None self.offset: int = 0 self.step: int = step def update_and_fetch(self, k_new: mx.array, v_new: mx.array) -> tuple[mx.array, mx.array]: """Append ``k_new``/``v_new`` to the slot and return the cumulative views. ``k_new``, ``v_new`` shape: (B, H, T_new, D). """ prev = self.offset new_len = prev + k_new.shape[2] if self.keys is None or new_len > self.keys.shape[2]: B, H, _, D = k_new.shape grow_to = self.step * ((new_len + self.step - 1) // self.step) grown_k = mx.zeros((B, H, grow_to, D), dtype=k_new.dtype) grown_v = mx.zeros((B, H, grow_to, D), dtype=v_new.dtype) if self.keys is not None: grown_k[..., :prev, :] = self.keys[..., :prev, :] grown_v[..., :prev, :] = self.values[..., :prev, :] self.keys = grown_k self.values = grown_v self.keys[..., prev:new_len, :] = k_new self.values[..., prev:new_len, :] = v_new self.offset = new_len return self.keys[..., :new_len, :], self.values[..., :new_len, :] class HrmKVCache: """128-slot recurrent cache (one per cycle-step × layer).""" def __init__(self, cfg: HrmTextConfig, step: int = 256): self.cfg = cfg self.slots = [_KVSlot(step) for _ in range(cfg.num_cache_slots)] @property def offset(self) -> int: """Number of sequence positions cached (each slot advances in lockstep).""" return self.slots[0].offset def slot(self, idx: int) -> _KVSlot: return self.slots[idx] # --------------------------------------------------------------------------- # Building blocks # --------------------------------------------------------------------------- class HrmAttention(nn.Module): """MHA + sigmoid output gate + RoPE. Fused gqkv projection: [gate | q | k | v].""" def __init__(self, hidden_size: int, num_heads: int, head_dim: int, rope_theta: float): super().__init__() self.num_heads = num_heads self.head_dim = head_dim self.scale = 1.0 / math.sqrt(head_dim) self.rope_theta = rope_theta self.gqkv_proj = nn.Linear(hidden_size, 4 * num_heads * head_dim, bias=False) self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False) def __call__( self, x: mx.array, mask: Optional[Union[str, mx.array]], rope_offset: int, slot: Optional[_KVSlot] = None, ) -> mx.array: B, S, _ = x.shape H, D = self.num_heads, self.head_dim gqkv = self.gqkv_proj(x) gate, q, k, v = mx.split(gqkv, 4, axis=-1) q = q.reshape(B, S, H, D).transpose(0, 2, 1, 3) k = k.reshape(B, S, H, D).transpose(0, 2, 1, 3) v = v.reshape(B, S, H, D).transpose(0, 2, 1, 3) gate = gate.reshape(B, S, H, D) # RoPE applied to q, k only — Llama / GPT-NeoX rotate_half convention. q = mxf.rope(q, D, traditional=False, base=self.rope_theta, scale=1.0, offset=rope_offset) k = mxf.rope(k, D, traditional=False, base=self.rope_theta, scale=1.0, offset=rope_offset) if slot is not None: k, v = slot.update_and_fetch(k, v) out = mxf.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask) out = out.transpose(0, 2, 1, 3) # (B, S, H, D) out = mx.sigmoid(gate) * out # output gating out = out.reshape(B, S, H * D) return self.o_proj(out) class HrmMLP(nn.Module): """SwiGLU with fused gate_up projection: [gate | up].""" def __init__(self, hidden_size: int, intermediate_size: int): super().__init__() self.gate_up_proj = nn.Linear(hidden_size, 2 * intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) def __call__(self, x: mx.array) -> mx.array: gu = self.gate_up_proj(x) gate, up = mx.split(gu, 2, axis=-1) return self.down_proj(nn.silu(gate) * up) class HrmLayer(nn.Module): """Pre-norm decoder block with parameterless RMSNorm.""" def __init__(self, cfg: HrmTextConfig): super().__init__() self.attn = HrmAttention(cfg.hidden_size, cfg.num_attention_heads, cfg.head_dim, cfg.rope_theta) self.mlp = HrmMLP(cfg.hidden_size, cfg.intermediate_size) self.eps = cfg.rms_norm_eps def __call__( self, x: mx.array, mask: Optional[Union[str, mx.array]], rope_offset: int, slot: Optional[_KVSlot], ) -> mx.array: x = x + self.attn(mxf.rms_norm(x, weight=None, eps=self.eps), mask, rope_offset, slot) x = x + self.mlp(mxf.rms_norm(x, weight=None, eps=self.eps)) return x class HrmStack(nn.Module): """One transformer stack — used twice as L_module and H_module.""" def __init__(self, cfg: HrmTextConfig): super().__init__() self.layers = [HrmLayer(cfg) for _ in range(cfg.num_layers_per_stack)] self.eps = cfg.rms_norm_eps def __call__( self, x: mx.array, mask: Optional[Union[str, mx.array]], rope_offset: int, cache: Optional[HrmKVCache], slot_base: int, ) -> mx.array: for i, layer in enumerate(self.layers): slot = cache.slot(slot_base + i) if cache is not None else None x = layer(x, mask, rope_offset, slot) return mxf.rms_norm(x, weight=None, eps=self.eps) # --------------------------------------------------------------------------- # Top-level model # --------------------------------------------------------------------------- class HrmTextModel(nn.Module): def __init__(self, cfg: HrmTextConfig): super().__init__() self.cfg = cfg self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) self.L_module = HrmStack(cfg) self.H_module = HrmStack(cfg) def make_cache(self, step: int = 256) -> HrmKVCache: return HrmKVCache(self.cfg, step=step) def __call__( self, input_ids: mx.array, cache: Optional[HrmKVCache] = None, token_type_ids: Optional[mx.array] = None, ) -> mx.array: B, S = input_ids.shape rope_offset = cache.offset if cache is not None else 0 z_H = self.embed_tokens(input_ids) * self.cfg.embedding_scale z_L = mx.zeros_like(z_H) mask = self._build_mask(S, rope_offset, token_type_ids, dtype=z_H.dtype) num_per_stack = self.cfg.num_layers_per_stack L_plus_1 = self.cfg.L_cycles + 1 for h in range(self.cfg.H_cycles): for l in range(self.cfg.L_cycles): slot_base = (h * L_plus_1 + l) * num_per_stack z_L = self.L_module(z_L + z_H, mask, rope_offset, cache, slot_base) slot_base = (h * L_plus_1 + self.cfg.L_cycles) * num_per_stack z_H = self.H_module(z_H + z_L, mask, rope_offset, cache, slot_base) return z_H def _build_mask( self, S: int, rope_offset: int, token_type_ids: Optional[mx.array], dtype, ) -> Optional[Union[str, mx.array]]: """Returns attention mask for fast.scaled_dot_product_attention. Fast path: prompt is all-prefix (token_type_ids=ones) on first call → full attention (return None). Single-token step (S==1) → return None (causal trivially). General path: build explicit additive mask of shape (1, 1, S, S+offset). """ if S == 1: return None if token_type_ids is None or self._is_all_prefix(token_type_ids): # First call, fully bidirectional within prefix block. return None # General PrefixLM mask: positions i in prefix and j in prefix attend mutually. # Otherwise causal (j <= i). idx = mx.arange(S) causal = idx[:, None] >= idx[None, :] tti = token_type_ids[0] == 1 prefix = tti[:, None] & tti[None, :] allow = causal | prefix m = mx.where(allow, mx.array(0.0, dtype=dtype), mx.array(-mx.inf, dtype=dtype)) return m[None, None, :, :] @staticmethod def _is_all_prefix(token_type_ids: mx.array) -> bool: # token_type_ids is (B, S) of int. Check first batch's positions are all 1. # We avoid synchronization by using a cheap reduction. return bool(mx.all(token_type_ids == 1).item()) class HrmTextForCausalLM(nn.Module): def __init__(self, cfg: HrmTextConfig): super().__init__() self.cfg = cfg self.model = HrmTextModel(cfg) self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) def __call__( self, input_ids: mx.array, cache: Optional[HrmKVCache] = None, token_type_ids: Optional[mx.array] = None, ) -> mx.array: h = self.model(input_ids, cache=cache, token_type_ids=token_type_ids) return self.lm_head(h) def make_cache(self, step: int = 256) -> HrmKVCache: return self.model.make_cache(step=step) # --------------------------------------------------------------------------- # Weight loading # --------------------------------------------------------------------------- def _load_safetensors_to_mlx(path: Path, dtype) -> dict: """Read a safetensors file containing bf16 tensors into MLX arrays. `safetensors` numpy/mlx backends can't decode bf16 directly, so we go through torch as a transient bridge. The bf16 → fp32 → mlx → bf16 round trip preserves values exactly for bf16 tensors (no precision loss). """ import torch out = {} with safe_open(str(path), framework="pt") as f: for k in f.keys(): t = f.get_tensor(k) out[k] = mx.array(t.float().numpy()).astype(dtype) return out def _load_mlx_native_safetensors(path: Path) -> dict: """Read an MLX-native safetensors file (already-mlx-saved).""" return mx.load(str(path)) def _apply_weights(model: HrmTextForCausalLM, params: dict) -> None: # z_L_init is zeros and frozen; we instantiate at call time. params = {k: v for k, v in params.items() if k != "model.z_L_init"} model.load_weights(list(params.items()), strict=True) def load(model_dir: Union[str, Path], dtype=mx.bfloat16) -> tuple[HrmTextForCausalLM, "Tokenizer"]: """Load HRM-Text-1B from a directory. The directory may contain either: - HF native safetensors (`model.safetensors` with torch bf16), or - MLX-converted safetensors (mlx.save_safetensors). Returns (model, tokenizer). """ model_dir = Path(model_dir) cfg_dict = json.loads((model_dir / "config.json").read_text()) cfg = HrmTextConfig.from_dict(cfg_dict) model = HrmTextForCausalLM(cfg) mlx_native = model_dir / "model_mlx.safetensors" hf_native = model_dir / "model.safetensors" if mlx_native.exists(): params = _load_mlx_native_safetensors(mlx_native) # MLX-native already filtered keys; trust them. model.load_weights(list(params.items()), strict=True) elif hf_native.exists(): params = _load_safetensors_to_mlx(hf_native, dtype) _apply_weights(model, params) else: raise FileNotFoundError(f"No safetensors found in {model_dir}") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(str(model_dir)) return model, tokenizer def convert_to_mlx(src_dir: Union[str, Path], dst_dir: Union[str, Path], dtype=mx.bfloat16) -> None: """One-shot conversion: HF safetensors → MLX-native safetensors. After conversion, `load(dst_dir)` reads the MLX file directly (no torch needed). """ src_dir = Path(src_dir) dst_dir = Path(dst_dir) dst_dir.mkdir(parents=True, exist_ok=True) params = _load_safetensors_to_mlx(src_dir / "model.safetensors", dtype) params = {k: v for k, v in params.items() if k != "model.z_L_init"} mx.save_safetensors(str(dst_dir / "model_mlx.safetensors"), params) # Copy config, tokenizer files import shutil for name in ["config.json", "tokenizer.json", "tokenizer_config.json", "LICENSE", "README.md"]: src = src_dir / name if src.exists(): shutil.copy(src, dst_dir / name) # --------------------------------------------------------------------------- # Sampling # --------------------------------------------------------------------------- def _sample(logits: mx.array, temperature: float, top_p: float, rng) -> int: """Sample one token id from (V,) logits.""" if temperature == 0.0: return int(mx.argmax(logits).item()) logits = logits.astype(mx.float32) / temperature if top_p < 1.0: sorted_idx = mx.argsort(-logits) sorted_logits = mx.take(logits, sorted_idx) sorted_probs = mx.softmax(sorted_logits, axis=-1) cumprobs = mx.cumsum(sorted_probs, axis=-1) # Keep tokens up to and including the first that crosses top_p. keep = cumprobs <= top_p # Always keep at least the top token. keep = mx.concatenate([mx.array([True]), keep[:-1]]) sorted_logits = mx.where(keep, sorted_logits, mx.array(-mx.inf)) # Restore original order. masked = mx.zeros_like(logits) + (-mx.inf) masked[sorted_idx] = sorted_logits logits = masked probs = mx.softmax(logits, axis=-1) # Categorical sampling token = int(mx.random.categorical(logits[None, :]).item()) return token # --------------------------------------------------------------------------- # Condition tag handling # --------------------------------------------------------------------------- _CONDITION_TOKEN = { "direct": "<|object_ref_start|>", "cot": "<|object_ref_end|>", "noisy": "<|quad_start|>", "synth": "<|quad_end|>", } def format_prompt(text: str, condition: str = "synth,cot") -> str: """Wrap a prompt with HRM-Text's <|im_start|> ... <|im_end|> envelope and the given composite condition prefix. `condition` is a comma-separated list of tags, e.g. "synth,cot" for the reasoning-style composite or "direct" for direct-answer mode. """ parts = [t.strip() for t in condition.split(",")] prefix = "".join(_CONDITION_TOKEN[p] for p in parts) return f"<|im_start|>{prefix}{text}<|im_end|>" # --------------------------------------------------------------------------- # Generation # --------------------------------------------------------------------------- def generate_step( model: HrmTextForCausalLM, prompt_ids: list[int], temperature: float = 0.0, top_p: float = 1.0, seed: int = 0, ) -> Iterator[int]: """Yields token ids one at a time. First yield is the first generated token (after prompt). Caller decides when to stop (EOS, max_tokens, etc.) by breaking out of the iterator. """ if temperature > 0.0: mx.random.seed(seed) cfg = model.cfg cache = model.make_cache() # Prefill: pass the whole prompt with token_type_ids=ones → bidirectional prefix. iid = mx.array([prompt_ids], dtype=mx.int32) tti = mx.ones((1, len(prompt_ids)), dtype=mx.int32) logits = model(iid, cache=cache, token_type_ids=tti) last = logits[0, -1, :] token = _sample(last, temperature, top_p, None) yield token while True: iid = mx.array([[token]], dtype=mx.int32) # No token_type_ids on subsequent steps → causal for new tokens. logits = model(iid, cache=cache) last = logits[0, -1, :] token = _sample(last, temperature, top_p, None) yield token def generate( model: HrmTextForCausalLM, tokenizer, prompt: str, condition: str = "synth,cot", max_tokens: int = 512, temperature: float = 0.0, top_p: float = 1.0, seed: int = 0, stream: bool = False, skip_special_tokens: bool = True, ) -> Union[str, Iterator[str]]: """High-level generation API. Args: condition: comma-separated condition tags. Common values: "synth,cot" - reasoning / chain-of-thought (DEFAULT, good for math) "synth" - synthetic / curated explanation "direct" - direct answer, no CoT (good for multi-choice, extraction) "cot" - chain-of-thought alone "noisy" - noisy / web-crawl style stream: if True returns an iterator that yields decoded text chunks. Returns: If stream=False: the decoded completion (str). If stream=True: iterator of decoded text fragments. """ formatted = format_prompt(prompt, condition) prompt_ids = tokenizer.encode(formatted, add_special_tokens=False) eos = model.cfg.eos_token_id def _gen_ids(): for i, tok in enumerate(generate_step(model, prompt_ids, temperature, top_p, seed)): if i >= max_tokens: return if eos is not None and tok == eos: return yield tok if stream: def _stream(): buf = [] for tok in _gen_ids(): buf.append(tok) # Decode incrementally — sometimes a token mid-multibyte produces "" which is fine. text = tokenizer.decode([tok], skip_special_tokens=skip_special_tokens) yield text return _stream() ids = list(_gen_ids()) return tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)