|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import json |
| 4 | +import re |
4 | 5 |
|
5 | 6 | from pathlib import Path |
6 | 7 | from typing import Callable, Iterable, TYPE_CHECKING |
@@ -355,3 +356,105 @@ def set_gguf_parameters(self): |
355 | 356 | self.gguf_writer.add_context_length(ctx_len) |
356 | 357 |
|
357 | 358 | self.gguf_writer.add_rope_dimension_sections(list(self.rope_parameters["xdrope_section"])) |
| 359 | + |
| 360 | + |
| 361 | +@ModelBase.register("HYV3ForCausalLM") |
| 362 | +class HYV3Model(TextModel): |
| 363 | + model_arch = gguf.MODEL_ARCH.HY_V3 |
| 364 | + |
| 365 | + # Trunk layer count, stashed before indexing so the classmethod |
| 366 | + # filter_tensors can identify the appended MTP block(s) (mirrors |
| 367 | + # Step35Model). |
| 368 | + _n_main_layers: int | None = None |
| 369 | + |
| 370 | + def __init__(self, *args, **kwargs): |
| 371 | + super().__init__(*args, **kwargs) |
| 372 | + # NextN/MTP layers are appended past num_hidden_layers; extend the |
| 373 | + # tensor map so the MTP block's tensors resolve to blk.<n>.* names. |
| 374 | + n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0)) |
| 375 | + if n_nextn > 0 and not self.no_mtp: |
| 376 | + self.block_count += n_nextn |
| 377 | + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) |
| 378 | + |
| 379 | + def index_tensors(self, remote_hf_model_id: str | None = None): |
| 380 | + type(self)._n_main_layers = self.hparams["num_hidden_layers"] |
| 381 | + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) |
| 382 | + |
| 383 | + def set_vocab(self): |
| 384 | + self._set_vocab_gpt2() |
| 385 | + |
| 386 | + def set_gguf_parameters(self): |
| 387 | + super().set_gguf_parameters() |
| 388 | + self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"]) |
| 389 | + self.gguf_writer.add_expert_shared_feed_forward_length( |
| 390 | + self.hparams["moe_intermediate_size"] * self.hparams.get("num_shared_experts", 1) |
| 391 | + ) |
| 392 | + self.gguf_writer.add_expert_weights_norm(self.hparams.get("route_norm", True)) |
| 393 | + self.gguf_writer.add_expert_weights_scale(float(self.hparams.get("router_scaling_factor", 1.0))) |
| 394 | + # sigmoid router with expert selection bias |
| 395 | + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) |
| 396 | + |
| 397 | + n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0)) |
| 398 | + if n_nextn > 0 and not self.no_mtp: |
| 399 | + self.gguf_writer.add_nextn_predict_layers(n_nextn) |
| 400 | + |
| 401 | + @classmethod |
| 402 | + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: |
| 403 | + if (titem := super().filter_tensors(item)) is None: |
| 404 | + return None |
| 405 | + name, gen = titem |
| 406 | + |
| 407 | + # HY V3 appends the MTP block(s) past num_hidden_layers. |
| 408 | + assert cls._n_main_layers is not None |
| 409 | + is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers |
| 410 | + |
| 411 | + # --no-mtp: drop the appended MTP block(s) entirely. |
| 412 | + if is_mtp and cls.no_mtp: |
| 413 | + return None |
| 414 | + # --mtp: keep ONLY MTP-block tensors plus the shared embeddings/norm/ |
| 415 | + # lm_head (so the resulting GGUF carries just the draft head). |
| 416 | + if cls.mtp_only and not is_mtp and name not in ( |
| 417 | + "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", |
| 418 | + ): |
| 419 | + return None |
| 420 | + |
| 421 | + # The MTP block's trailing final_layernorm (applied after the decoder |
| 422 | + # block, before the shared LM head) maps to nextn.shared_head_norm. |
| 423 | + if is_mtp: |
| 424 | + name = name.replace(".final_layernorm.", ".shared_head.norm.") |
| 425 | + |
| 426 | + return name, gen |
| 427 | + |
| 428 | + _experts: list[dict[str, Tensor]] | None = None |
| 429 | + |
| 430 | + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: |
| 431 | + # merge the per-expert tensors into stacked 3d tensors |
| 432 | + if name.startswith("model.layers.") and ".mlp.experts." in name: |
| 433 | + n_experts = self.find_hparam(["num_local_experts", "num_experts"]) |
| 434 | + assert bid is not None |
| 435 | + |
| 436 | + if self._experts is None: |
| 437 | + self._experts = [{} for _ in range(self.block_count)] |
| 438 | + |
| 439 | + self._experts[bid][name] = data_torch |
| 440 | + |
| 441 | + if len(self._experts[bid]) >= n_experts * 3: |
| 442 | + for w_name in ("down_proj", "gate_proj", "up_proj"): |
| 443 | + datas: list[Tensor] = [] |
| 444 | + for xid in range(n_experts): |
| 445 | + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" |
| 446 | + datas.append(self._experts[bid][ename]) |
| 447 | + del self._experts[bid][ename] |
| 448 | + |
| 449 | + merged = torch.stack(datas, dim=0) |
| 450 | + yield from super().modify_tensors(merged, f"model.layers.{bid}.mlp.experts.{w_name}.weight", bid) |
| 451 | + return |
| 452 | + |
| 453 | + yield from super().modify_tensors(data_torch, name, bid) |
| 454 | + |
| 455 | + def prepare_tensors(self): |
| 456 | + super().prepare_tensors() |
| 457 | + if self._experts is not None: |
| 458 | + experts = [k for d in self._experts for k in d.keys()] |
| 459 | + if experts: |
| 460 | + raise ValueError(f"Unprocessed experts: {experts}") |
0 commit comments