char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenconvert: add @ModelBase.example (#27208)fa88ae936

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.2 KiB80 linesraw
1from __future__ import annotations
2
3from typing import Callable, Iterable, TYPE_CHECKING
4
5import torch
6
7if TYPE_CHECKING:
8    from torch import Tensor
9
10from .base import ModelBase, gguf
11
12from .llama import LlamaModel
13
14
15@ModelBase.register("AfmoeForCausalLM")
16@ModelBase.example("arcee-ai/Trinity-Large-Thinking")
17class AfmoeModel(LlamaModel):
18    model_arch = gguf.MODEL_ARCH.AFMOE
19
20    def set_gguf_parameters(self):
21        super().set_gguf_parameters()
22
23        # MoE parameters
24        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 scaling
32        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 attention
38        if (sliding_window := self.hparams.get("sliding_window")) is not None:
39            self.gguf_writer.add_sliding_window(sliding_window)
40
41    @classmethod
42    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
43        name, gen = item
44
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 format
52        # process the experts separately
53        if name.find("mlp.experts") != -1:
54            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
55            assert bid is not None
56
57            if self._experts is None:
58                self._experts = [{} for _ in range(self.block_count)]
59
60            self._experts[bid][name] = data_torch
61
62            if len(self._experts[bid]) >= n_experts * 3:
63                # merge the experts into a single 3d tensor
64                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                return
77            else:
78                return
79
80        yield from ModelBase.modify_tensors(self, data_torch, name, bid)