diff --git a/conversion/__init__.py b/conversion/__init__.py index 9c5d984388..d48861e46b 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -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", diff --git a/conversion/base.py b/conversion/base.py index d2d80be368..8f6b3519cb 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -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" diff --git a/conversion/hrm_text.py b/conversion/hrm_text.py new file mode 100644 index 0000000000..3684fe912a --- /dev/null +++ b/conversion/hrm_text.py @@ -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) diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index 6af74cd874..3a15a6fca3 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -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"}, ] diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index e54ee5a0fd..36b4c3190e 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -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, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index ed5a185b32..56cc65a70a 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -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) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 0fac27efc0..03b7951c3c 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -135,6 +135,7 @@ static const std::map 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_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_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_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}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 6e67f5d659..af64f4ec57 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -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, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 6334f3ccab..f21767601d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -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(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(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 || diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 3afa49ebe8..73dffcc9f7 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -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; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 59a8ff84f8..0a27367c98 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -11,6 +11,7 @@ #include #include +#include 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 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(&layer)[i]); + const struct ggml_tensor * tensor = reinterpret_cast(&layer)[i]; + if (tensor == nullptr || !seen.insert(tensor).second) { + continue; + } + add_tensor(tensor); } } } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3b2536283c..3148f27816 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -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: diff --git a/src/llama-model.h b/src/llama-model.h index a02b30ca7a..d61afa2cb3 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -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; diff --git a/src/models/hrm-text.cpp b/src/models/hrm-text.cpp new file mode 100644 index 0000000000..4a9a67b6c1 --- /dev/null +++ b/src/models/hrm-text.cpp @@ -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 llama_model_hrm_text::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*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); +} diff --git a/src/models/models.h b/src/models/models.h index da519dcfdc..3f9c67c63c 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -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 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; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 90a6a71623..568f7234c1 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -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); }