mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-17 20:31:47 +02:00
model : add support for HrmTextForCausalLM (DFM Mimir 1B) (#27625)
* model : add support for HrmTextForCausalLM (DFM Mimir 1B) HRM-Text runs two transformer stacks (low, high) in an alternating cycle over the same token stream. The low-cycle state z_l starts from a learned [n_embd] tensor and is broadcast over positions. - conversion: new writer for the fused gqkv projection (order gate,q,k,v) remapped to llama.cpp q/k/v plus a separate sigmoid gate tensor - loader: block_count = lps * h_cycles * (l_cycles + 1) cache slots aliasing 2*lps physical blocks via struct copies - graph: looped build with sigmoid-gated attention, SwiGLU FFN and parameterless RMS norms; learned embedding_scale applied in build_inp_embd - saver: pointer-deduplicated layer loop (looped archs alias tensors) - tests: hrm_text fixture (lps 1, h 2, l 3) in test-llama-archs Limitations: causal attention only - the upstream prefix-LM mode is not implemented (the prefix_lm GGUF key round-trips unused). The KV cache holds one entry per pass: 128 layers for Mimir 1B, i.e. 4x a same-width 32-layer model - about 3072 MiB at ctx 4096 in F16 (halves with q8_0 KV + FA). Every token runs all 128 block passes, so decode cost is roughly 4x a dense model of equal width (2.65 t/s BF16, 8-thread desktop CPU). Verified against the HF reference: identical argmax at 334/334 positions across 20 prompts (BF16 GGUF vs FP32 golden). q8_0 requant: 95.8% top-1, all remaining misses inside the HF top-5 (accumulated error over 128 sequential blocks). AI usage disclosure: YES Used GLM-5.3 for the majority of code AI-generated under my direction, all gates verified locally. All in all I could say that I have written less than 20% of the code and most of the heavy lifting has been done by the model. As such, this should be considered experimental. * Update conversion/hrm_text.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update src/llama-arch.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * convert : add gguf_writer methods for hrm_text metadata replace raw add_uint32/add_bool calls with dedicated GGUFWriter methods, following the add_embedding_scale pattern Assisted-by: GLM-5.3 * convert : map regular hrm_text tensors via tensor_mapping delegate unfused checkpoints to the base tensor mapping; training-style attn. names are renamed to self_attn. so the patterns match Assisted-by: GLM-5.3 * model : format hrm-text build_* calls as in other models one argument group per line, matching sibling model files Assisted-by: GLM-5.3 * llama : move hrm z_l_init table entries out of the nemotron group place the name and tensor-info entries with the other global input tensors Assisted-by: GLM-5.3 * convert : slim down hrm_text comments Assisted-by: GLM-5.3 * convert : build hrm_text block tensor names from the {bid} template The tensor map holds concrete per-block names, so format the template with the computed layer index before handing it to super(). * llama : name hrm metadata keys in their own hrm. namespace The four keys are arch-independent, unlike the arch-substituted Keys.LLM entries, so group them under Keys.HRM (like Keys.Split) and rename the llm_kv entries to LLM_KV_HRM_*. Only our own GGUFs carry the old hrm_text.* keys; they are regenerated. * Update src/llama-model-saver.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * llama : keep hrm metadata keys arch-substituted Per review: the GGUF keys stay "{arch}.h_cycles" style, so the Python members drop the LLM_KV_HRM_ prefix and keep arch templates; C++ keeps the LLM_KV_HRM_* enums. GGUF output is unchanged - existing files and HF uploads stay valid. * Update gguf-py/gguf/constants.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update src/llama-arch.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update src/llama-arch.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * convert : rename hrm writer methods to add_hrm_* Generic names like add_h_cycles/add_prefix_lm are too broad on the shared GGUFWriter; prefix them with hrm_ like the metadata keys. * model : fix meta-split lookup for archs with aliased cache slots Cache tensors of archs that alias physical blocks across looped slots (hrm_text, nanbeige with num_loops > 1) can reference block indices without weight tensor names. Take the output projection from the layer array instead of asserting; all other lookups are unchanged. * model : replicate hrm_text tensors on meta devices instead of splitting The aliased cache slots rotate split states differently from their physical weights, so the meta-split execution invariants (set_rows requires the cache state to match the token indices) cannot hold for any device count. Replicate all hrm_text tensors on every meta device instead; single-device and non-meta paths are unchanged. Assisted-by: Claude Sonnet --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
This commit is contained in:
co-authored by
Sigbjørn Skjæret
parent
83078fec0d
commit
7d6f5d02bb
@@ -123,6 +123,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"HunYuanDenseV1ForCausalLM": "hunyuan",
|
||||
"HunYuanMoEV1ForCausalLM": "hunyuan",
|
||||
"HunYuanVLForConditionalGeneration": "hunyuan",
|
||||
"HrmTextForCausalLM": "hrm_text",
|
||||
"HYV3ForCausalLM": "hunyuan",
|
||||
"HYV4ForCausalLM": "hy_v4",
|
||||
"IQuestCoderForCausalLM": "llama",
|
||||
|
||||
@@ -1633,6 +1633,9 @@ class TextModel(ModelBase):
|
||||
if chkhsh == "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7":
|
||||
# ref: https://huggingface.co/LiquidAI/LFM2.5-8B-A1B
|
||||
res = "lfm2"
|
||||
if chkhsh == "846deafc5b0fa786186fa4ae6c7b49903cf2f1d1895bdb80b9120d60be135252":
|
||||
# ref: https://huggingface.co/danish-foundation-models/DFM-Mimir
|
||||
res = "gemma4"
|
||||
if chkhsh == "0a766d034107bc736a3f2dc4968fd62e54a3570f1454443e0c5a4cc6bd7941ed":
|
||||
# ref: https://huggingface.co/XHToken/Spark-X2.5-1.7B
|
||||
res = "spark2_5"
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from typing import Iterable, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch import Tensor
|
||||
|
||||
from .base import ModelBase, TextModel, gguf
|
||||
|
||||
|
||||
@ModelBase.register("HrmTextForCausalLM")
|
||||
@ModelBase.example("danish-foundation-models/DFM-Mimir")
|
||||
class HrmTextModel(TextModel):
|
||||
model_arch = gguf.MODEL_ARCH.HRM_TEXT
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
# training-style configs store the per-stack count in num_hidden_layers,
|
||||
# transformers-style configs keep it in num_layers_per_stack
|
||||
self.layers_per_stack = self.hparams.get("num_layers_per_stack") or self.hparams["num_hidden_layers"]
|
||||
self.h_cycles = self.hparams["H_cycles"]
|
||||
self.l_cycles = self.hparams["L_cycles"]
|
||||
|
||||
# block_count is the expanded cache-slot count; the file only holds
|
||||
# 2 * layers_per_stack physical blocks
|
||||
self.block_count = self.layers_per_stack * self.h_cycles * (self.l_cycles + 1)
|
||||
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, 2 * self.layers_per_stack)
|
||||
|
||||
def set_vocab(self):
|
||||
self._set_vocab_gpt2()
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
|
||||
self.gguf_writer.add_rope_dimension_count(head_dim)
|
||||
self.gguf_writer.add_embedding_scale(self.hparams["embedding_scale"])
|
||||
self.gguf_writer.add_hrm_layers_per_stack(self.layers_per_stack)
|
||||
self.gguf_writer.add_hrm_h_cycles(self.h_cycles)
|
||||
self.gguf_writer.add_hrm_l_cycles(self.l_cycles)
|
||||
self.gguf_writer.add_hrm_prefix_lm(bool(self.hparams.get("prefix_lm", False)))
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
if name == "model.embed_tokens.weight":
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), data_torch
|
||||
return
|
||||
if name == "lm_head.weight":
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT), data_torch
|
||||
return
|
||||
if name == "model.z_L_init":
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.HRM_Z_L_INIT, suffix=""), data_torch
|
||||
return
|
||||
|
||||
match = re.fullmatch(r"model\.([LH])_module\.layers\.(\d+)\.(.+)", name)
|
||||
if match is None:
|
||||
raise ValueError(f"can not map tensor: {name}")
|
||||
|
||||
stack, layer_s, tensor_name = match.groups()
|
||||
# the L stack occupies blocks [0, layers_per_stack), the H stack follows it
|
||||
layer_idx = int(layer_s) + (self.layers_per_stack if stack == "H" else 0)
|
||||
|
||||
if tensor_name == "attn.gqkv_proj.weight":
|
||||
gate, q, k, v = data_torch.chunk(4, dim=0)
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_GATE, layer_idx), gate.contiguous()
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, layer_idx), q.contiguous()
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, layer_idx), k.contiguous()
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, layer_idx), v.contiguous()
|
||||
elif tensor_name == "mlp.gate_up_proj.weight":
|
||||
gate, up = data_torch.chunk(2, dim=0)
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, layer_idx), gate.contiguous()
|
||||
yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, layer_idx), up.contiguous()
|
||||
else:
|
||||
if tensor_name.startswith("attn."):
|
||||
tensor_name = "self_attn." + tensor_name[len("attn."):]
|
||||
tensor_name = "model.layers.{bid}." + tensor_name
|
||||
yield from super().modify_tensors(data_torch, tensor_name.format(bid=layer_idx), layer_idx)
|
||||
@@ -191,6 +191,10 @@ pre_computed_hashes = [
|
||||
{"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/evilfreelancer/ruGPT3XL", "chkhsh": "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4"},
|
||||
# lfm2 variants
|
||||
{"name": "lfm2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/LiquidAI/LFM2.5-8B-A1B", "chkhsh": "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7"},
|
||||
# hrm-text (DFM Mimir) is SPM-style BPE: normalizer maps ' ' -> '▁', merges
|
||||
# over the whole text (fix_mistral_regex inserts a tekken regex that is a
|
||||
# no-op here); the gemma4 pre (escape ws, split on newlines only) matches it.
|
||||
{"name": "gemma4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/danish-foundation-models/DFM-Mimir", "chkhsh": "846deafc5b0fa786186fa4ae6c7b49903cf2f1d1895bdb80b9120d60be135252"},
|
||||
{"name": "spark2_5", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/XHToken/Spark-X2.5-1.7B", "chkhsh": "0a766d034107bc736a3f2dc4968fd62e54a3570f1454443e0c5a4cc6bd7941ed"},
|
||||
]
|
||||
|
||||
|
||||
@@ -277,6 +277,12 @@ class Keys:
|
||||
LLM_KV_SPLIT_COUNT = "split.count"
|
||||
LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"
|
||||
|
||||
class HRM:
|
||||
LAYERS_PER_STACK = "{arch}.hrm.layers_per_stack"
|
||||
H_CYCLES = "{arch}.hrm.h_cycles"
|
||||
L_CYCLES = "{arch}.hrm.l_cycles"
|
||||
PREFIX_LM = "{arch}.hrm.prefix_lm"
|
||||
|
||||
class SSM:
|
||||
CONV_KERNEL = "{arch}.ssm.conv_kernel"
|
||||
INNER_SIZE = "{arch}.ssm.inner_size"
|
||||
@@ -511,6 +517,7 @@ class MODEL_ARCH(IntEnum):
|
||||
QWEN3 = auto()
|
||||
QWEN3MOE = auto()
|
||||
QWEN3NEXT = auto()
|
||||
HRM_TEXT = auto()
|
||||
QWEN3VL = auto()
|
||||
QWEN3VLMOE = auto()
|
||||
QWEN35 = auto()
|
||||
@@ -655,6 +662,7 @@ class MODEL_TENSOR(IntEnum):
|
||||
TOKEN_TYPES = auto()
|
||||
POS_EMBD = auto()
|
||||
OUTPUT = auto()
|
||||
HRM_Z_L_INIT = auto()
|
||||
DENSE_2_OUT = auto() # embeddinggemma 2_Dense
|
||||
DENSE_3_OUT = auto() # embeddinggemma 3_Dense
|
||||
OUTPUT_NORM = auto()
|
||||
@@ -1266,6 +1274,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.QWEN3: "qwen3",
|
||||
MODEL_ARCH.QWEN3MOE: "qwen3moe",
|
||||
MODEL_ARCH.QWEN3NEXT: "qwen3next",
|
||||
MODEL_ARCH.HRM_TEXT: "hrm_text",
|
||||
MODEL_ARCH.QWEN3VL: "qwen3vl",
|
||||
MODEL_ARCH.QWEN3VLMOE: "qwen3vlmoe",
|
||||
MODEL_ARCH.QWEN35: "qwen35",
|
||||
@@ -1410,6 +1419,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.POS_EMBD: "position_embd",
|
||||
MODEL_TENSOR.OUTPUT_NORM: "output_norm",
|
||||
MODEL_TENSOR.OUTPUT: "output",
|
||||
MODEL_TENSOR.HRM_Z_L_INIT: "hrm.z_l_init",
|
||||
MODEL_TENSOR.DENSE_2_OUT: "dense_2", # embeddinggemma 2_Dense
|
||||
MODEL_TENSOR.DENSE_3_OUT: "dense_3", # embeddinggemma 2_Dense
|
||||
MODEL_TENSOR.HC_HEAD_FN: "output_hc_fn",
|
||||
@@ -2796,6 +2806,19 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
|
||||
MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
|
||||
],
|
||||
MODEL_ARCH.HRM_TEXT: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
MODEL_TENSOR.HRM_Z_L_INIT,
|
||||
MODEL_TENSOR.ATTN_Q,
|
||||
MODEL_TENSOR.ATTN_K,
|
||||
MODEL_TENSOR.ATTN_V,
|
||||
MODEL_TENSOR.ATTN_GATE,
|
||||
MODEL_TENSOR.ATTN_OUT,
|
||||
MODEL_TENSOR.FFN_GATE,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
],
|
||||
MODEL_ARCH.QWEN3VL: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
|
||||
@@ -932,6 +932,18 @@ class GGUFWriter:
|
||||
def add_embedding_scale(self, value: float) -> None:
|
||||
self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value)
|
||||
|
||||
def add_hrm_layers_per_stack(self, value: int) -> None:
|
||||
self.add_uint32(Keys.HRM.LAYERS_PER_STACK.format(arch=self.arch), value)
|
||||
|
||||
def add_hrm_h_cycles(self, value: int) -> None:
|
||||
self.add_uint32(Keys.HRM.H_CYCLES.format(arch=self.arch), value)
|
||||
|
||||
def add_hrm_l_cycles(self, value: int) -> None:
|
||||
self.add_uint32(Keys.HRM.L_CYCLES.format(arch=self.arch), value)
|
||||
|
||||
def add_hrm_prefix_lm(self, value: bool) -> None:
|
||||
self.add_bool(Keys.HRM.PREFIX_LM.format(arch=self.arch), value)
|
||||
|
||||
def add_adapter_count(self, count: int) -> None:
|
||||
self.add_uint32(Keys.Adapters.COUNT.format(arch=self.arch), count)
|
||||
|
||||
|
||||
@@ -135,6 +135,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_GROVEMOE, "grovemoe" },
|
||||
{ LLM_ARCH_APERTUS, "apertus" },
|
||||
{ LLM_ARCH_MINIMAX_01, "minimax-01" },
|
||||
{ LLM_ARCH_HRM_TEXT, "hrm_text" },
|
||||
{ LLM_ARCH_MINIMAX_M2, "minimax-m2" },
|
||||
{ LLM_ARCH_MINIMAX_M3, "minimax-m3" },
|
||||
{ LLM_ARCH_COGVLM, "cogvlm" },
|
||||
@@ -246,6 +247,10 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" },
|
||||
{ LLM_KV_NUM_LOOPS, "%s.num_loops" },
|
||||
{ LLM_KV_SKIP_LOOP_FINAL_NORM, "%s.skip_loop_final_norm" },
|
||||
{ LLM_KV_HRM_LAYERS_PER_STACK, "%s.hrm.layers_per_stack" },
|
||||
{ LLM_KV_HRM_H_CYCLES, "%s.hrm.h_cycles" },
|
||||
{ LLM_KV_HRM_L_CYCLES, "%s.hrm.l_cycles" },
|
||||
{ LLM_KV_HRM_PREFIX_LM, "%s.hrm.prefix_lm" },
|
||||
|
||||
{ LLM_KV_ATTENTION_HEAD_COUNT, "%s.attention.head_count" },
|
||||
{ LLM_KV_ATTENTION_HEAD_COUNT_KV, "%s.attention.head_count_kv" },
|
||||
@@ -431,6 +436,7 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
|
||||
{ LLM_TENSOR_OUTPUT_NORM_LFM2, "token_embd_norm" }, // fix for wrong tensor name
|
||||
{ LLM_TENSOR_OUTPUT, "output" },
|
||||
{ LLM_TENSOR_ROPE_FREQS, "rope_freqs" },
|
||||
{ LLM_TENSOR_HRM_Z_L_INIT, "hrm.z_l_init" },
|
||||
{ LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" },
|
||||
{ LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" },
|
||||
{ LLM_TENSOR_ATTN_K, "blk.%d.attn_k" },
|
||||
@@ -714,6 +720,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
|
||||
{LLM_TENSOR_TOKEN_EMBD, {LLM_TENSOR_LAYER_INPUT, GGML_OP_GET_ROWS}},
|
||||
{LLM_TENSOR_POS_EMBD, {LLM_TENSOR_LAYER_INPUT, GGML_OP_GET_ROWS}},
|
||||
{LLM_TENSOR_TOKEN_TYPES, {LLM_TENSOR_LAYER_INPUT, GGML_OP_GET_ROWS}},
|
||||
{LLM_TENSOR_HRM_Z_L_INIT, {LLM_TENSOR_LAYER_INPUT, GGML_OP_ADD}},
|
||||
{LLM_TENSOR_TOKEN_EMBD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, // do the norms on the first layer (not the input layer)
|
||||
{LLM_TENSOR_OUTPUT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_CLS, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
|
||||
|
||||
@@ -162,6 +162,7 @@ enum llm_arch {
|
||||
LLM_ARCH_QWEN3TTS,
|
||||
LLM_ARCH_POCKETTTS,
|
||||
LLM_ARCH_MINIMAX_01,
|
||||
LLM_ARCH_HRM_TEXT,
|
||||
LLM_ARCH_UNKNOWN,
|
||||
};
|
||||
|
||||
@@ -251,6 +252,10 @@ enum llm_kv {
|
||||
LLM_KV_FULL_ATTENTION_INTERVAL,
|
||||
LLM_KV_NUM_LOOPS,
|
||||
LLM_KV_SKIP_LOOP_FINAL_NORM,
|
||||
LLM_KV_HRM_LAYERS_PER_STACK,
|
||||
LLM_KV_HRM_H_CYCLES,
|
||||
LLM_KV_HRM_L_CYCLES,
|
||||
LLM_KV_HRM_PREFIX_LM,
|
||||
|
||||
LLM_KV_ATTENTION_HEAD_COUNT,
|
||||
LLM_KV_ATTENTION_HEAD_COUNT_KV,
|
||||
@@ -693,6 +698,7 @@ enum llm_tensor {
|
||||
LLM_TENSOR_NEXTN_SHARED_HEAD_NORM,
|
||||
LLM_TENSOR_MASKED_EMBD_CENTROIDS,
|
||||
LLM_TENSOR_MASKED_EMBD_ORDERING,
|
||||
LLM_TENSOR_HRM_Z_L_INIT,
|
||||
LLM_TENSOR_FC,
|
||||
LLM_TENSOR_D2T,
|
||||
LLM_TENSOR_DSPARK_MARKOV_W1,
|
||||
|
||||
@@ -2309,6 +2309,9 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
|
||||
if (model.arch == LLM_ARCH_KIMI_K3) {
|
||||
// the n_tokens*40 budget below is exhausted at ubatch 3840
|
||||
res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors());
|
||||
} else if (model.arch == LLM_ARCH_HRM_TEXT) {
|
||||
// the 128-slot looped graph needs roughly one stack per token budget
|
||||
res = std::max<uint32_t>(n_tokens * 80, 64u * model.n_tensors());
|
||||
} else if (model.arch == LLM_ARCH_QWEN3NEXT ||
|
||||
model.arch == LLM_ARCH_KIMI_LINEAR ||
|
||||
model.arch == LLM_ARCH_BAILINGMOE3 ||
|
||||
|
||||
@@ -206,6 +206,12 @@ struct llama_hparams {
|
||||
float situ_beta = 1.0f;
|
||||
float situ_linear_beta = 0.0f; // 0 = no linear-beta transform on the up branch
|
||||
|
||||
// hrm-text (looped H/L stacks)
|
||||
uint32_t n_hrm_layers_per_stack = 0;
|
||||
uint32_t n_hrm_h_cycles = 0;
|
||||
uint32_t n_hrm_l_cycles = 0;
|
||||
bool hrm_prefix_lm = false;
|
||||
|
||||
bool ssm_dt_b_c_rms = false;
|
||||
|
||||
float f_clamp_kqv = 0.0f;
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <unordered_set>
|
||||
|
||||
bool llama_model_saver_supports_arch(llm_arch arch) {
|
||||
switch (arch) {
|
||||
@@ -261,6 +262,10 @@ void llama_model_saver::add_kv_from_model() {
|
||||
add_kv(LLM_KV_TIME_DECAY_EXTRA_DIM, hparams.time_decay_extra_dim);
|
||||
add_kv(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale);
|
||||
add_kv(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale);
|
||||
add_kv(LLM_KV_HRM_LAYERS_PER_STACK, hparams.n_hrm_layers_per_stack);
|
||||
add_kv(LLM_KV_HRM_H_CYCLES, hparams.n_hrm_h_cycles);
|
||||
add_kv(LLM_KV_HRM_L_CYCLES, hparams.n_hrm_l_cycles);
|
||||
add_kv(LLM_KV_HRM_PREFIX_LM, hparams.hrm_prefix_lm);
|
||||
add_kv(LLM_KV_TOKEN_SHIFT_COUNT, hparams.token_shift_count);
|
||||
add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, hparams.n_moe_layer_step);
|
||||
// add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, ???); // saved as LLM_KV_ATTENTION_RECURRENT_LAYERS instead
|
||||
@@ -475,6 +480,7 @@ void llama_model_saver::add_tensors_from_model() {
|
||||
add_tensor(model->cls_out);
|
||||
add_tensor(model->cls_out_b);
|
||||
add_tensor(model->cls_norm);
|
||||
add_tensor(model->hrm_z_l_init);
|
||||
add_tensor(model->hc_head_fn);
|
||||
add_tensor(model->hc_head_base);
|
||||
add_tensor(model->hc_head_scale);
|
||||
@@ -483,9 +489,17 @@ void llama_model_saver::add_tensors_from_model() {
|
||||
add_tensor(model->hc_head_down);
|
||||
add_tensor(model->hc_head_up);
|
||||
|
||||
// looped architectures alias physical tensors across cache slots; save each
|
||||
// tensor once. a different tensor with an existing name still asserts below
|
||||
std::unordered_set<const struct ggml_tensor *> seen;
|
||||
|
||||
for (const struct llama_layer & layer : model->layers) {
|
||||
for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) {
|
||||
add_tensor(reinterpret_cast<const struct ggml_tensor * const *>(&layer)[i]);
|
||||
const struct ggml_tensor * tensor = reinterpret_cast<const struct ggml_tensor * const *>(&layer)[i];
|
||||
if (tensor == nullptr || !seen.insert(tensor).second) {
|
||||
continue;
|
||||
}
|
||||
add_tensor(tensor);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -314,6 +314,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
|
||||
return new llama_model_minimax_m2(params);
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
return new llama_model_minimax_m3(params);
|
||||
case LLM_ARCH_HRM_TEXT:
|
||||
return new llama_model_hrm_text(params);
|
||||
case LLM_ARCH_COGVLM:
|
||||
return new llama_model_cogvlm(params);
|
||||
case LLM_ARCH_PANGU_EMBED:
|
||||
@@ -473,6 +475,10 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
|
||||
};
|
||||
|
||||
auto get_tensor_config = [&]() -> tensor_config {
|
||||
if (ud->model->arch == LLM_ARCH_HRM_TEXT) {
|
||||
// aliased cache slots cannot satisfy the meta-split invariants, so replicate all tensors
|
||||
return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, tensor, 0, 0};
|
||||
}
|
||||
if (is_dsv4) {
|
||||
if (std::regex_match(tensor_name, pattern_kv_cache) ||
|
||||
std::regex_match(tensor_name, pattern_dsv4_state)) {
|
||||
@@ -3022,6 +3028,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
|
||||
case LLM_ARCH_TALKIE:
|
||||
case LLM_ARCH_MELLUM:
|
||||
case LLM_ARCH_MAPLE:
|
||||
case LLM_ARCH_HRM_TEXT:
|
||||
return LLAMA_ROPE_TYPE_NEOX;
|
||||
|
||||
case LLM_ARCH_DFLASH:
|
||||
|
||||
@@ -643,6 +643,9 @@ struct llama_model {
|
||||
struct ggml_tensor * nextn_proj_pre = nullptr;
|
||||
struct ggml_tensor * nextn_proj_post = nullptr;
|
||||
|
||||
// hrm-text initial low-cycle state
|
||||
struct ggml_tensor * hrm_z_l_init = nullptr;
|
||||
|
||||
// DeepSeek-V4
|
||||
struct ggml_tensor * hc_head_fn = nullptr;
|
||||
struct ggml_tensor * hc_head_base = nullptr;
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
#include "models.h"
|
||||
|
||||
// HRM-Text: alternating low/high transformer stacks over the same token stream.
|
||||
// Reference: HrmTextModel in transformers, DFM Mimir 1B.
|
||||
|
||||
void llama_model_hrm_text::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
||||
ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false);
|
||||
|
||||
ml.get_key(LLM_KV_HRM_LAYERS_PER_STACK, hparams.n_hrm_layers_per_stack);
|
||||
ml.get_key(LLM_KV_HRM_H_CYCLES, hparams.n_hrm_h_cycles);
|
||||
ml.get_key(LLM_KV_HRM_L_CYCLES, hparams.n_hrm_l_cycles);
|
||||
|
||||
// prefix-LM prefill is not implemented (causal attention only); kept for round-trip
|
||||
ml.get_key(LLM_KV_HRM_PREFIX_LM, hparams.hrm_prefix_lm, false);
|
||||
|
||||
GGML_ASSERT(hparams.n_hrm_layers_per_stack > 0);
|
||||
GGML_ASSERT(hparams.n_hrm_h_cycles > 0);
|
||||
GGML_ASSERT(hparams.n_hrm_l_cycles > 0);
|
||||
|
||||
// the GGUF block count is the expanded cache-slot count
|
||||
const uint32_t n_slot = hparams.n_hrm_layers_per_stack * hparams.n_hrm_h_cycles * (hparams.n_hrm_l_cycles + 1);
|
||||
GGML_ASSERT(hparams.n_layer() == n_slot);
|
||||
|
||||
switch (hparams.n_embd) {
|
||||
case 1536:
|
||||
type = LLM_TYPE_1B;
|
||||
break;
|
||||
default:
|
||||
type = LLM_TYPE_UNKNOWN;
|
||||
}
|
||||
}
|
||||
|
||||
void llama_model_hrm_text::load_arch_tensors(llama_model_loader &) {
|
||||
LLAMA_LOAD_LOCALS;
|
||||
|
||||
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
|
||||
|
||||
// output
|
||||
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
|
||||
// if output is NULL, init from the input tok embed
|
||||
if (output == NULL) {
|
||||
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
|
||||
}
|
||||
|
||||
hrm_z_l_init = create_tensor(tn(LLM_TENSOR_HRM_Z_L_INIT), { n_embd }, 0);
|
||||
|
||||
const int lps = hparams.n_hrm_layers_per_stack;
|
||||
|
||||
// blocks [0, lps) hold the low stack, blocks [lps, 2*lps) hold the high stack.
|
||||
// the first low and high passes create the layers; later passes alias them.
|
||||
const int l_first = 0;
|
||||
const int h_first = hparams.n_hrm_l_cycles * lps;
|
||||
|
||||
for (int h = 0; h < (int) hparams.n_hrm_h_cycles; ++h) {
|
||||
for (int l = 0; l < (int) hparams.n_hrm_l_cycles + 1; ++l) {
|
||||
const int slot_base = (h * (hparams.n_hrm_l_cycles + 1) + l) * lps;
|
||||
const int blk_base = l == (int) hparams.n_hrm_l_cycles ? lps : 0;
|
||||
|
||||
if (h > 0 || (l > 0 && l < (int) hparams.n_hrm_l_cycles)) {
|
||||
// alias pass: these cache slots hold the same layers as the first passes
|
||||
const int src_base = l == (int) hparams.n_hrm_l_cycles ? h_first : l_first;
|
||||
for (int il = 0; il < lps; ++il) {
|
||||
layers[slot_base + il] = layers[src_base + il];
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
for (int il = 0; il < lps; ++il) {
|
||||
auto & layer = layers[slot_base + il];
|
||||
const int bid = blk_base + il;
|
||||
|
||||
create_tensor_qkv(layer, bid, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
|
||||
|
||||
// sigmoid attention gate, applied to the attention output before o_proj
|
||||
layer.wqkv_gate =
|
||||
create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", bid), { n_embd, n_embd_head_k * n_head }, 0);
|
||||
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", bid), { n_embd_head_k * n_head, n_embd }, 0);
|
||||
|
||||
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", bid), { n_embd, n_ff }, 0);
|
||||
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", bid), { n_ff, n_embd }, 0);
|
||||
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", bid), { n_embd, n_ff }, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<llm_graph_context> llama_model_hrm_text::build_arch_graph(const llm_graph_params & params) const {
|
||||
return std::make_unique<graph>(*this, params);
|
||||
}
|
||||
|
||||
// one stack invocation: lps pre-norm decoder layers, then the parameterless final norm
|
||||
ggml_tensor * llama_model_hrm_text::graph::build_stack(llm_graph_input_attn_kv * inp_attn,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * cur,
|
||||
int slot_base) const {
|
||||
const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
|
||||
|
||||
const int lps = model.hparams.n_hrm_layers_per_stack;
|
||||
|
||||
for (int il = 0; il < lps; ++il) {
|
||||
const int s = slot_base + il;
|
||||
const auto & layer = model.layers[s];
|
||||
|
||||
ggml_tensor * inpSA = cur;
|
||||
|
||||
cur = build_norm(cur, nullptr, nullptr, LLM_NORM_RMS, s);
|
||||
cb(cur, "attn_norm", s);
|
||||
|
||||
// sigmoid-gated self-attention (same shape as qwen3next attention layers)
|
||||
{
|
||||
ggml_tensor * gate = build_lora_mm(layer.wqkv_gate, cur);
|
||||
cb(gate, "attn_gate_proj", s);
|
||||
|
||||
auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head_k, n_head, n_head_kv, s);
|
||||
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
cb(Qcur, "Qcur", s);
|
||||
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
cb(Kcur, "Kcur", s);
|
||||
|
||||
cur = build_attn(inp_attn,
|
||||
nullptr, nullptr, nullptr,
|
||||
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, s);
|
||||
cb(cur, "attn_pregate", s);
|
||||
|
||||
gate = ggml_sigmoid(ctx0, gate);
|
||||
cb(gate, "attn_gate_sigmoid", s);
|
||||
|
||||
cur = ggml_mul(ctx0, cur, gate);
|
||||
cb(cur, "attn_gated", s);
|
||||
|
||||
cur = build_lora_mm(layer.wo, cur, layer.wo_s);
|
||||
cb(cur, "attn_out", s);
|
||||
}
|
||||
|
||||
cur = ggml_add(ctx0, cur, inpSA);
|
||||
cb(cur, "attn_add", s);
|
||||
|
||||
inpSA = cur;
|
||||
cur = build_norm(cur, nullptr, nullptr, LLM_NORM_RMS, s);
|
||||
cb(cur, "ffn_norm", s);
|
||||
|
||||
cur = build_ffn(cur,
|
||||
layer.ffn_up, nullptr, nullptr,
|
||||
layer.ffn_gate, nullptr, nullptr,
|
||||
layer.ffn_down, nullptr, nullptr,
|
||||
nullptr,
|
||||
LLM_FFN_SILU, LLM_FFN_PAR, s);
|
||||
cb(cur, "ffn_out", s);
|
||||
|
||||
cur = ggml_add(ctx0, cur, inpSA);
|
||||
cb(cur, "ffn_add", s);
|
||||
|
||||
cur = build_cvec(cur, s);
|
||||
cb(cur, "l_out", s);
|
||||
}
|
||||
|
||||
cur = build_norm(cur, nullptr, nullptr, LLM_NORM_RMS, slot_base);
|
||||
cb(cur, "stack_norm", slot_base);
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
llama_model_hrm_text::graph::graph(const llama_model & model, const llm_graph_params & params) :
|
||||
llm_graph_context(params),
|
||||
model(model) {
|
||||
ggml_tensor * cur;
|
||||
|
||||
// {n_embd, n_tokens}, scaled by hparams.f_embedding_scale inside build_inp_embd
|
||||
ggml_tensor * zH = build_inp_embd(model.tok_embd);
|
||||
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
|
||||
auto * inp_attn = build_attn_inp_kv();
|
||||
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
// the learned low-cycle state is [n_embd]; binary ops broadcast it over [n_embd, n_tokens]
|
||||
ggml_tensor * zL = model.hrm_z_l_init;
|
||||
|
||||
for (uint32_t h = 0; h < model.hparams.n_hrm_h_cycles; ++h) {
|
||||
for (uint32_t l = 0; l < model.hparams.n_hrm_l_cycles; ++l) {
|
||||
const int slot_base = (h * (model.hparams.n_hrm_l_cycles + 1) + l) * model.hparams.n_hrm_layers_per_stack;
|
||||
|
||||
zL = build_stack(inp_attn, inp_pos, ggml_add(ctx0, zH, zL), slot_base);
|
||||
}
|
||||
|
||||
const int slot_base = (h * (model.hparams.n_hrm_l_cycles + 1) + model.hparams.n_hrm_l_cycles) *
|
||||
model.hparams.n_hrm_layers_per_stack;
|
||||
|
||||
zH = build_stack(inp_attn, inp_pos, ggml_add(ctx0, zH, zL), slot_base);
|
||||
}
|
||||
|
||||
cur = zH;
|
||||
|
||||
if (inp_out_ids) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
}
|
||||
|
||||
cb(cur, "result_norm", -1);
|
||||
res->t_embd = cur;
|
||||
|
||||
cur = build_lora_mm(model.output, cur, model.output_s);
|
||||
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
@@ -1825,6 +1825,27 @@ struct llama_model_plm : public llama_model_base {
|
||||
};
|
||||
|
||||
|
||||
struct llama_model_hrm_text : public llama_model_base {
|
||||
llama_model_hrm_text(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
void load_arch_tensors(llama_model_loader & ml) override;
|
||||
|
||||
struct graph : public llm_graph_context {
|
||||
graph(const llama_model & model, const llm_graph_params & params);
|
||||
|
||||
const llama_model & model;
|
||||
|
||||
ggml_tensor * build_stack(
|
||||
llm_graph_input_attn_kv * inp_attn,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * cur,
|
||||
int slot_base) const;
|
||||
};
|
||||
|
||||
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
|
||||
};
|
||||
|
||||
|
||||
struct llama_model_bailingmoe : public llama_model_base {
|
||||
llama_model_bailingmoe(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
|
||||
@@ -130,6 +130,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
} else if (arch == LLM_ARCH_QWEN3TTS) {
|
||||
//n_vocab = 4096; // must be >= the hard-coded codec head size (3072)
|
||||
n_vocab = 3072; // TODO: should be 4096, but user code cannot get `n_vocab_out` yet [TAG_LLAMA_N_VOCAB_OUT]
|
||||
} else if (arch == LLM_ARCH_HRM_TEXT) {
|
||||
n_layer = 8; // 1 layer per stack x 2 h-cycles x (3 l-cycles + 1) cache slots
|
||||
}
|
||||
|
||||
uint32_t n_head_kv = n_head;
|
||||
@@ -325,6 +327,13 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true);
|
||||
}
|
||||
|
||||
if (arch == LLM_ARCH_HRM_TEXT) {
|
||||
// 8 cache slots alias 2 physical blocks: 1 low-stack layer + 1 high-stack layer
|
||||
ms.add_kv(LLM_KV_HRM_LAYERS_PER_STACK, uint32_t(1));
|
||||
ms.add_kv(LLM_KV_HRM_H_CYCLES, uint32_t(2));
|
||||
ms.add_kv(LLM_KV_HRM_L_CYCLES, uint32_t(3));
|
||||
}
|
||||
|
||||
if (arch == LLM_ARCH_MAPLE) {
|
||||
ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 7.0f);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user