Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf, logger11 12 13@ModelBase.register("GroveMoeForCausalLM", "modeling_grove_moe.GroveMoeForCausalLM")14@ModelBase.example("inclusionAI/GroveMoE-Inst")15class GroveMoeModel(TextModel):16    model_arch = gguf.MODEL_ARCH.GROVEMOE17 18    def set_gguf_parameters(self):19        super().set_gguf_parameters()20        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:21            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)22            logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}")23        # FIXME?: Hardcoded https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L29924        self.gguf_writer.add_expert_chunk_feed_forward_length(self.hparams.get("head_dim") or 128)25        # FIXME?: Hardcoded https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L29826        self.gguf_writer.add_experts_per_group(2)27        # FIXME?: Hardcoded https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L37628        self.gguf_writer.add_expert_group_scale(0.05)29 30    _experts: list[dict[str, Tensor]] | None = None31    _chunk_experts: list[dict[str, Tensor]] | None = None32 33    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:34        if name.endswith(".expert_bias"):35            # FIXME?: Unused https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L30336            return37 38        # process the experts separately39        if name.find("chunk_experts") != -1:40            n_experts = self.find_hparam(["num_local_experts", "num_experts"]) // 2 # see add_experts_per_group41            assert bid is not None42 43            if self._chunk_experts is None:44                self._chunk_experts = [{} for _ in range(self.block_count)]45 46            self._chunk_experts[bid][name] = data_torch47 48            if len(self._chunk_experts[bid]) >= n_experts * 3:49                # merge the experts into a single 3d tensor50                for w_name in ["down_proj", "gate_proj", "up_proj"]:51                    datas: list[Tensor] = []52 53                    for xid in range(n_experts):54                        ename = f"model.layers.{bid}.mlp.chunk_experts.{xid}.{w_name}.weight"55                        datas.append(self._chunk_experts[bid][ename])56                        del self._chunk_experts[bid][ename]57 58                    data_torch = torch.stack(datas, dim=0)59 60                    merged_name = f"model.layers.{bid}.mlp.chunk_experts.{w_name}.weight"61 62                    yield from super().modify_tensors(data_torch, merged_name, bid)63                return64            else:65                return66        elif name.find("experts") != -1:67            n_experts = self.find_hparam(["num_local_experts", "num_experts"])68            assert bid is not None69 70            if self._experts is None:71                self._experts = [{} for _ in range(self.block_count)]72 73            self._experts[bid][name] = data_torch74 75            if len(self._experts[bid]) >= n_experts * 3:76                # merge the experts into a single 3d tensor77                for w_name in ["down_proj", "gate_proj", "up_proj"]:78                    datas: list[Tensor] = []79 80                    for xid in range(n_experts):81                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"82                        datas.append(self._experts[bid][ename])83                        del self._experts[bid][ename]84 85                    data_torch = torch.stack(datas, dim=0)86 87                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"88 89                    yield from super().modify_tensors(data_torch, merged_name, bid)90                return91            else:92                return93 94        yield from super().modify_tensors(data_torch, name, bid)95 96    def prepare_tensors(self):97        super().prepare_tensors()98 99        if self._chunk_experts is not None:100            # flatten `list[dict[str, Tensor]]` into `list[str]`101            chunk_experts = [k for d in self._chunk_experts for k in d.keys()]102            if len(chunk_experts) > 0:103                raise ValueError(f"Unprocessed adjugate experts: {chunk_experts}")104 105        if self._experts is not None:106            # flatten `list[dict[str, Tensor]]` into `list[str]`107            experts = [k for d in self._experts for k in d.keys()]108            if len(experts) > 0:109                raise ValueError(f"Unprocessed experts: {experts}")110