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 →
2.5 KiB62 linesraw
1from __future__ import annotations
2
3from typing import Iterable, TYPE_CHECKING
4
5import torch
6
7if TYPE_CHECKING:
8    from torch import Tensor
9
10from .base import ModelBase, TextModel, gguf, logger
11
12
13@ModelBase.register("MellumForCausalLM")
14@ModelBase.example("JetBrains/Mellum2-12B-A2.5B-Base")
15class MellumModel(TextModel):
16    model_arch = gguf.MODEL_ARCH.MELLUM
17
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
24        use_sliding_window = self.hparams.get("use_sliding_window")
25        sliding_window = self.hparams.get("sliding_window")
26        if (use_sliding_window is True or use_sliding_window is None) and sliding_window is not None:
27            self.gguf_writer.add_sliding_window(sliding_window)
28            logger.info(f"gguf: sliding window = {sliding_window}")
29            self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]])
30            logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}")
31
32    _experts: list[dict[str, Tensor]] | None = None
33
34    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
35        if name.find("experts") != -1:
36            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
37            assert bid is not None
38
39            if self._experts is None:
40                self._experts = [{} for _ in range(self.block_count)]
41
42            self._experts[bid][name] = data_torch
43
44            if len(self._experts[bid]) >= n_experts * 3:
45                for w_name in ["down_proj", "gate_proj", "up_proj"]:
46                    datas: list[Tensor] = []
47
48                    for xid in range(n_experts):
49                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
50                        datas.append(self._experts[bid][ename])
51                        del self._experts[bid][ename]
52
53                    data_torch = torch.stack(datas, dim=0)
54
55                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
56
57                    yield from super().modify_tensors(data_torch, merged_name, bid)
58                return
59            else:
60                return
61
62        yield from super().modify_tensors(data_torch, name, bid)