Buckets:
| 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 | |
| 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) | |
Xet Storage Details
- Size:
- 3.8 kB
- Xet hash:
- 6ef49f7a06b10c8e9f99667fd829c890ac7c423093a997de9f6a415a28c0773f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.