Buckets:
| from __future__ import annotations | |
| from typing import Callable, Iterable, TYPE_CHECKING | |
| import torch | |
| if TYPE_CHECKING: | |
| from torch import Tensor | |
| from .base import ModelBase, TextModel, gguf, logger | |
| class MellumModel(TextModel): | |
| model_arch = gguf.MODEL_ARCH.MELLUM | |
| supports_mtp_export = True | |
| _n_main_layers: int | None = None | |
| _experts: list[dict[str, Tensor]] | None = None | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| # number of exported MTP blocks, appended after the trunk layers | |
| self.n_nextn = 0 if self.no_mtp else (self.find_hparam(["num_nextn_predict_layers", "num_mtp_layers"], optional=True) or 0) | |
| if self.mtp_only and self.n_nextn == 0: | |
| raise ValueError("--mtp requested, but the checkpoint has no MTP layers (num_nextn_predict_layers is 0)") | |
| if self.n_nextn > 0: | |
| if any(t != "sparse" for t in self.hparams.get("mtp_mlp_layer_types") or []): | |
| raise ValueError("Mellum MTP layers with dense MLP are not supported") | |
| self.block_count += self.n_nextn | |
| self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) | |
| def index_tensors(self, remote_hf_model_id: str | None = None): | |
| type(self)._n_main_layers = self.hparams["num_hidden_layers"] | |
| return super().index_tensors(remote_hf_model_id=remote_hf_model_id) | |
| def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: | |
| if (titem := super().filter_tensors(item)) is None: | |
| return None | |
| name, gen = titem | |
| if name.startswith("mtp."): | |
| if cls.no_mtp: | |
| return None | |
| assert cls._n_main_layers is not None | |
| # mtp.layers.{i}.* -> model.layers.{n_main + i}.*, mtp.{fc,norm,...} -> nextn names of the first MTP block | |
| remapper = { | |
| "fc": "eh_proj", | |
| "pre_fc_norm_embedding": "enorm", | |
| "pre_fc_norm_hidden": "hnorm", | |
| "norm": "shared_head.norm", | |
| } | |
| parts = name.split(".", 3) | |
| if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal(): | |
| name = f"model.layers.{cls._n_main_layers + int(parts[2])}.{parts[3]}" | |
| elif len(parts) == 3 and parts[1] in remapper: | |
| name = f"model.layers.{cls._n_main_layers}.{remapper[parts[1]]}.{parts[2]}" | |
| elif cls.mtp_only and name not in ("model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"): | |
| return None | |
| return name, gen | |
| def set_gguf_parameters(self): | |
| super().set_gguf_parameters() | |
| if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None: | |
| self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size) | |
| logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}") | |
| use_sliding_window = self.hparams.get("use_sliding_window") | |
| sliding_window = self.hparams.get("sliding_window") | |
| if (use_sliding_window is True or use_sliding_window is None) and sliding_window is not None: | |
| self.gguf_writer.add_sliding_window(sliding_window) | |
| logger.info(f"gguf: sliding window = {sliding_window}") | |
| layer_types = list(self.hparams["layer_types"]) | |
| # MTP blocks default to sliding window attention, as in transformers | |
| if self.n_nextn > 0: | |
| layer_types += self.hparams.get("mtp_layer_types") or ["sliding_attention"] * self.n_nextn | |
| self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in layer_types]) | |
| logger.info(f"gguf: sliding window pattern length = {len(layer_types)}") | |
| if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None: | |
| self.gguf_writer.add_expert_weights_norm(norm_topk_prob) | |
| if self.n_nextn > 0: | |
| self.gguf_writer.add_nextn_predict_layers(self.n_nextn) | |
| def prepare_metadata(self, vocab_only: bool): | |
| from_dir = self.fname_out.is_dir() | |
| super().prepare_metadata(vocab_only=vocab_only) | |
| if not self.mtp_only or not from_dir: | |
| return | |
| output_type: str = self.ftype.name.partition("_")[2] | |
| fname_default: str = gguf.naming_convention( | |
| self.metadata.name, self.metadata.basename, self.metadata.finetune, | |
| self.metadata.version, size_label=None, output_type=output_type, model_type=None) | |
| self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" | |
| def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: | |
| if name.find("experts") != -1: | |
| n_experts = self.find_hparam(["num_local_experts", "num_experts"]) | |
| assert bid is not None | |
| if self._experts is None: | |
| self._experts = [{} for _ in range(self.block_count)] | |
| self._experts[bid][name] = data_torch | |
| if len(self._experts[bid]) >= n_experts * 3: | |
| for w_name in ["down_proj", "gate_proj", "up_proj"]: | |
| datas: list[Tensor] = [] | |
| for xid in range(n_experts): | |
| ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" | |
| datas.append(self._experts[bid][ename]) | |
| del self._experts[bid][ename] | |
| data_torch = torch.stack(datas, dim=0) | |
| merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" | |
| yield from super().modify_tensors(data_torch, merged_name, bid) | |
| return | |
| else: | |
| return | |
| yield from super().modify_tensors(data_torch, name, bid) | |
Xet Storage Details
- Size:
- 5.99 kB
- Xet hash:
- 86aa6c8d1e0f566d853620d81d8890ee61b177cab1d5e195fa2c2344fed570ff
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.