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 →
8.6 KiB204 linesraw
1from __future__ import annotations
2
3import json
4import math
5import re
6
7from typing import Callable, Iterable, TYPE_CHECKING
8
9import torch
10
11if TYPE_CHECKING:
12    from torch import Tensor
13
14from .base import MmprojModel, ModelBase, TextModel, gguf
15
16
17@ModelBase.register("Ernie4_5_ForCausalLM", "Ernie4_5ForCausalLM")
18@ModelBase.example("baidu/ERNIE-4.5-0.3B-PT")
19class Ernie4_5Model(TextModel):
20    model_arch = gguf.MODEL_ARCH.ERNIE4_5
21
22    def set_vocab(self):
23        self._set_vocab_sentencepiece()
24
25        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'
26        if tokenizer_config_file.is_file():
27            with open(tokenizer_config_file, "r", encoding="utf-8") as f:
28                tokenizer_config_json = json.load(f)
29                if "add_prefix_space" in tokenizer_config_json:
30                    self.gguf_writer.add_add_space_prefix(tokenizer_config_json["add_prefix_space"])
31
32    def set_gguf_parameters(self):
33        super().set_gguf_parameters()
34
35    @classmethod
36    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
37        name, gen = item
38
39        if "ernie." in name:
40            name = name.replace("ernie.", "model.")
41
42        return super().filter_tensors((name, gen))
43
44    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
45        num_heads = self.hparams["num_attention_heads"]
46        num_kv_heads = self.hparams["num_key_value_heads"]
47        if (head_dim := self.hparams.get("head_dim")) is None:
48            head_dim = self.hparams["hidden_size"] // num_heads
49
50        # split the qkv weights
51        # qkv_proj shape: [(num_heads + 2 * num_kv_heads) * head_dim, hidden_size]
52        if "qkv_proj" in name:
53            name_q = name.replace("qkv_proj.weight", "q_proj.weight")
54            name_k = name.replace("qkv_proj.weight", "k_proj.weight")
55            name_v = name.replace("qkv_proj.weight", "v_proj.weight")
56            total_q_dim = num_heads * head_dim
57            total_k_dim = num_kv_heads * head_dim
58            total_v_dim = num_kv_heads * head_dim
59            q_proj_weight, k_proj_weight, v_proj_weight = data_torch.split([total_q_dim, total_k_dim, total_v_dim], dim=0)
60            yield from super().modify_tensors(q_proj_weight, name_q, bid)
61            yield from super().modify_tensors(k_proj_weight, name_k, bid)
62            yield from super().modify_tensors(v_proj_weight, name_v, bid)
63        # split the up_gate_proj into gate and up
64        # up_gate_proj shape: [2 * intermediate_size, hidden_size]
65        elif "up_gate_proj" in name:
66            name_up = name.replace("up_gate_proj.weight", "up_proj.weight")
67            name_gate = name.replace("up_gate_proj.weight", "gate_proj.weight")
68            dim_half = data_torch.shape[0] // 2
69            gate_proj_weight, up_proj_weight = data_torch.split(dim_half, dim=0)
70            yield from super().modify_tensors(gate_proj_weight, name_gate, bid)
71            yield from super().modify_tensors(up_proj_weight, name_up, bid)
72        else:
73            yield from super().modify_tensors(data_torch, name, bid)
74
75
76@ModelBase.register("Ernie4_5_MoeForCausalLM")
77@ModelBase.example("baidu/ERNIE-4.5-21B-A3B-PT")
78class Ernie4_5MoeModel(Ernie4_5Model):
79    model_arch = gguf.MODEL_ARCH.ERNIE4_5_MOE
80    _experts: list[dict[str, Tensor]] | None = None
81
82    def __init__(self, *args, **kwargs):
83        super().__init__(*args, **kwargs)
84        self._experts = [{} for _ in range(self.block_count)]
85
86    def set_gguf_parameters(self):
87        super().set_gguf_parameters()
88        self.gguf_writer.add_expert_count(self.hparams["moe_num_experts"])
89        self.gguf_writer.add_expert_used_count(self.hparams["moe_k"])
90        self.gguf_writer.add_interleave_moe_layer_step(self.hparams["moe_layer_interval"])
91        self.gguf_writer.add_leading_dense_block_count(self.hparams["moe_layer_start_index"])
92        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:
93            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)
94        if (shared_expert_count := self.hparams.get('moe_num_shared_experts')) is not None:
95            self.gguf_writer.add_expert_shared_count(shared_expert_count)
96            if shared_expert_count > 0 and (shared_expert_intermediate_size := self.hparams.get('intermediate_size')) is not None and (num_key_value_heads := self.hparams.get('num_key_value_heads')) is not None:
97                self.gguf_writer.add_expert_shared_feed_forward_length(shared_expert_intermediate_size // num_key_value_heads)
98
99    @classmethod
100    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
101        name, gen = item
102
103        # skip Multi-Token Prediction (MTP) layers (again, same as DeepseekV2)
104        match = re.match(r"model.mtp_block.(\d+)", name)
105        if match:
106            return None
107
108        # skip all other MTP tensors for now
109        match = re.match(r"model.mtp_emb_norm.(\d+)", name)
110        if match:
111            return None
112
113        match = re.match(r"model.mtp_hidden_norm.(\d+)", name)
114        if match:
115            return None
116
117        match = re.match(r"model.mtp_linear_proj.(\d+)", name)
118        if match:
119            return None
120
121        return super().filter_tensors(item)
122
123    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
124        # process the experts separately
125        if name.find("mlp.experts") != -1:
126            n_experts = self.hparams["moe_num_experts"]
127            assert bid is not None
128
129            if self._experts is None:
130                self._experts = [{} for _ in range(self.block_count)]
131
132            self._experts[bid][name] = data_torch
133
134            if len(self._experts[bid]) >= n_experts * 3:
135                # merge the experts into a single 3d tensor
136                for w_name in ["gate_proj", "up_proj", "down_proj"]:
137                    datas: list[Tensor] = []
138
139                    for xid in range(n_experts):
140                        ename_to_retrieve = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
141                        datas.append(self._experts[bid][ename_to_retrieve])
142                        del self._experts[bid][ename_to_retrieve]
143
144                    data_torch = torch.stack(datas, dim=0)
145                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
146                    yield from super().modify_tensors(data_torch, merged_name, bid)
147        else:
148            yield from ModelBase.modify_tensors(self, data_torch, name, bid)
149
150    def prepare_tensors(self):
151        super().prepare_tensors()
152
153        if self._experts is not None:
154            # flatten `list[dict[str, Tensor]]` into `list[str]`
155            experts = [k for d in self._experts for k in d.keys()]
156            if len(experts) > 0:
157                raise ValueError(f"Unprocessed experts: {experts}")
158
159
160@ModelBase.register("PaddleOCRVLForConditionalGeneration")
161@ModelBase.example("PaddlePaddle/PaddleOCR-VL")
162class PaddleOCRModel(Ernie4_5Model):
163    model_arch = gguf.MODEL_ARCH.PADDLEOCR
164
165
166@ModelBase.register("PaddleOCRVisionModel")
167@ModelBase.example("PaddlePaddle/PaddleOCR-VL")
168class PaddleOCRVisionModel(MmprojModel):
169    # PaddleOCR-VL uses a modified version of Siglip
170    min_pixels: int = 0
171    max_pixels: int = 0
172
173    def __init__(self, *args, **kwargs):
174        super().__init__(*args, **kwargs)
175        assert self.hparams_vision is not None
176        self.min_pixels = self.preprocessor_config["min_pixels"]
177        self.max_pixels = self.preprocessor_config["max_pixels"]
178        self.hparams_vision["image_size"] = int(math.sqrt(self.max_pixels))
179
180    def set_gguf_parameters(self):
181        super().set_gguf_parameters()
182        assert self.hparams_vision is not None
183        hparams = self.hparams_vision
184        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.PADDLEOCR)
185        self.gguf_writer.add_vision_max_pixels(self.max_pixels)
186        self.gguf_writer.add_vision_min_pixels(self.min_pixels)
187        self.gguf_writer.add_vision_use_gelu(True)
188        self.gguf_writer.add_vision_attention_layernorm_eps(hparams.get("rms_norm_eps", 1e-6))
189
190    @classmethod
191    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
192        name, gen = item
193
194        if "vision_model" not in name and "mlp_AR" not in name:
195            return None
196        name = name.replace("visual.", "model.")
197        if "packing_position_embedding" in name:
198            # unused
199            return None
200        if "vision_model.head" in name:
201            # we don't yet support image embeddings for this model
202            return None
203
204        return super().filter_tensors((name, gen))