Brunobkr/llama.cpp_AlgMor24_github

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")14class GroveMoeModel(TextModel):15    model_arch = gguf.MODEL_ARCH.GROVEMOE16 17    def set_gguf_parameters(self):18        super().set_gguf_parameters()19        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:20            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)21            logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}")22        # FIXME?: Hardcoded https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L29923        self.gguf_writer.add_expert_chunk_feed_forward_length(self.hparams.get("head_dim") or 128)24        # FIXME?: Hardcoded https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L29825        self.gguf_writer.add_experts_per_group(2)26        # FIXME?: Hardcoded https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L37627        self.gguf_writer.add_expert_group_scale(0.05)28 29    _experts: list[dict[str, Tensor]] | None = None30    _chunk_experts: list[dict[str, Tensor]] | None = None31 32    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:33        if name.endswith(".expert_bias"):34            # FIXME?: Unused https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L30335            return36 37        # process the experts separately38        if name.find("chunk_experts") != -1:39            n_experts = self.find_hparam(["num_local_experts", "num_experts"]) // 2 # see add_experts_per_group40            assert bid is not None41 42            if self._chunk_experts is None:43                self._chunk_experts = [{} for _ in range(self.block_count)]44 45            self._chunk_experts[bid][name] = data_torch46 47            if len(self._chunk_experts[bid]) >= n_experts * 3:48                # merge the experts into a single 3d tensor49                for w_name in ["down_proj", "gate_proj", "up_proj"]:50                    datas: list[Tensor] = []51 52                    for xid in range(n_experts):53                        ename = f"model.layers.{bid}.mlp.chunk_experts.{xid}.{w_name}.weight"54                        datas.append(self._chunk_experts[bid][ename])55                        del self._chunk_experts[bid][ename]56 57                    data_torch = torch.stack(datas, dim=0)58 59                    merged_name = f"model.layers.{bid}.mlp.chunk_experts.{w_name}.weight"60 61                    yield from super().modify_tensors(data_torch, merged_name, bid)62                return63            else:64                return65        elif name.find("experts") != -1:66            n_experts = self.find_hparam(["num_local_experts", "num_experts"])67            assert bid is not None68 69            if self._experts is None:70                self._experts = [{} for _ in range(self.block_count)]71 72            self._experts[bid][name] = data_torch73 74            if len(self._experts[bid]) >= n_experts * 3:75                # merge the experts into a single 3d tensor76                for w_name in ["down_proj", "gate_proj", "up_proj"]:77                    datas: list[Tensor] = []78 79                    for xid in range(n_experts):80                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"81                        datas.append(self._experts[bid][ename])82                        del self._experts[bid][ename]83 84                    data_torch = torch.stack(datas, dim=0)85 86                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"87 88                    yield from super().modify_tensors(data_torch, merged_name, bid)89                return90            else:91                return92 93        yield from super().modify_tensors(data_torch, name, bid)94 95    def prepare_tensors(self):96        super().prepare_tensors()97 98        if self._chunk_experts is not None:99            # flatten `list[dict[str, Tensor]]` into `list[str]`100            chunk_experts = [k for d in self._chunk_experts for k in d.keys()]101            if len(chunk_experts) > 0:102                raise ValueError(f"Unprocessed adjugate experts: {chunk_experts}")103 104        if self._experts is not None:105            # flatten `list[dict[str, Tensor]]` into `list[str]`106            experts = [k for d in self._experts for k in d.keys()]107            if len(experts) > 0:108                raise ValueError(f"Unprocessed experts: {experts}")109