Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Callable, Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, gguf11 12from .llama import LlamaModel13 14 15@ModelBase.register("AfmoeForCausalLM")16@ModelBase.example("arcee-ai/Trinity-Large-Thinking")17class AfmoeModel(LlamaModel):18    model_arch = gguf.MODEL_ARCH.AFMOE19 20    def set_gguf_parameters(self):21        super().set_gguf_parameters()22 23        # MoE parameters24        if (n_shared_experts := self.hparams.get("num_shared_experts")) is not None:25            self.gguf_writer.add_expert_shared_count(n_shared_experts)26        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:27            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)28        if (n_dense_layers := self.hparams.get("num_dense_layers")) is not None:29            self.gguf_writer.add_leading_dense_block_count(n_dense_layers)30 31        # Route normalization and scaling32        if (route_norm := self.hparams.get("route_norm")) is not None:33            self.gguf_writer.add_expert_weights_norm(route_norm)34        if (route_scale := self.hparams.get("route_scale")) is not None:35            self.gguf_writer.add_expert_weights_scale(route_scale)36 37        # Sliding window attention38        if (sliding_window := self.hparams.get("sliding_window")) is not None:39            self.gguf_writer.add_sliding_window(sliding_window)40 41    @classmethod42    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:43        name, gen = item44 45        if name.endswith(".expert_bias"):46            name = name.replace(".expert_bias", ".expert_bias.bias")47 48        return super().filter_tensors((name, gen))49 50    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:51        # Handle expert weights - they're already merged in the HF format52        # process the experts separately53        if name.find("mlp.experts") != -1:54            n_experts = self.find_hparam(["num_local_experts", "num_experts"])55            assert bid is not None56 57            if self._experts is None:58                self._experts = [{} for _ in range(self.block_count)]59 60            self._experts[bid][name] = data_torch61 62            if len(self._experts[bid]) >= n_experts * 3:63                # merge the experts into a single 3d tensor64                for w_name in ["gate_proj", "up_proj", "down_proj"]:65                    datas: list[Tensor] = []66 67                    for xid in range(n_experts):68                        ename_to_retrieve = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"69                        datas.append(self._experts[bid][ename_to_retrieve])70                        del self._experts[bid][ename_to_retrieve]71 72                    data_torch = torch.stack(datas, dim=0)73                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"74                    yield from ModelBase.modify_tensors(self, data_torch, merged_name, bid)75 76                return77            else:78                return79 80        yield from ModelBase.modify_tensors(self, data_torch, name, bid)81