80 lines
3.7 KiB
Python
80 lines
3.7 KiB
Python
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)
|