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 →
4.4 KiB124 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
11
12from .llama import LlamaModel
13
14
15@ModelBase.register("OlmoForCausalLM")
16@ModelBase.register("OLMoForCausalLM")
17@ModelBase.example("allenai/OLMo-1.7-7B-hf")
18class OlmoModel(TextModel):
19    model_arch = gguf.MODEL_ARCH.OLMO
20
21    def set_gguf_parameters(self):
22        super().set_gguf_parameters()
23        self.gguf_writer.add_layer_norm_eps(1e-5)
24        clip_qkv = self.hparams.get("clip_qkv")
25        if clip_qkv is not None:
26            self.gguf_writer.add_clamp_kqv(clip_qkv)
27
28    # Same as super class, but permuting q_proj, k_proj
29    # Copied from: LlamaModel
30    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
31        n_head = self.hparams["num_attention_heads"]
32        n_kv_head = self.hparams.get("num_key_value_heads")
33
34        if name.endswith("q_proj.weight"):
35            data_torch = LlamaModel.permute(data_torch, n_head, n_head)
36        if name.endswith("k_proj.weight"):
37            data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)
38
39        yield from super().modify_tensors(data_torch, name, bid)
40
41
42@ModelBase.register("SeedOssForCausalLM")
43@ModelBase.example("ByteDance-Seed/Seed-OSS-36B-Instruct")
44class SeedOssModel(TextModel):
45    model_arch = gguf.MODEL_ARCH.SEED_OSS
46
47
48@ModelBase.register("Olmo2ForCausalLM")
49@ModelBase.register("Olmo3ForCausalLM")
50@ModelBase.example("allenai/OLMo-2-1124-7B-Instruct", "allenai/Olmo-3-7B-Instruct")
51class Olmo2Model(TextModel):
52    model_arch = gguf.MODEL_ARCH.OLMO2
53
54    def set_gguf_parameters(self):
55        super().set_gguf_parameters()
56
57        if "sliding_window" in self.hparams:
58            self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
59
60            sliding_window_pattern = []
61            if "layer_types" in self.hparams:
62                sliding_window_pattern = [t == "sliding_attention" for t in self.hparams["layer_types"]]
63            else:
64                # Olmo2 does not use sliding window attention.
65                # Olmo3 defaults to using sliding window for all layers except every 4th.
66                for i in range(self.hparams["num_hidden_layers"]):
67                    sliding_window_pattern.append((i + 1) % 4 != 0)
68
69            self.gguf_writer.add_sliding_window_pattern(sliding_window_pattern)
70
71
72@ModelBase.register("OlmoeForCausalLM")
73@ModelBase.example("allenai/OLMoE-1B-7B-0924")
74class OlmoeModel(TextModel):
75    model_arch = gguf.MODEL_ARCH.OLMOE
76
77    def set_gguf_parameters(self):
78        super().set_gguf_parameters()
79        self.gguf_writer.add_layer_norm_rms_eps(1e-5)
80
81    _experts: list[dict[str, Tensor]] | None = None
82
83    # Copied from: Qwen2MoeModel
84    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
85        # process the experts separately
86        if name.find("experts") != -1:
87            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
88            assert bid is not None
89
90            if self._experts is None:
91                self._experts = [{} for _ in range(self.block_count)]
92
93            self._experts[bid][name] = data_torch
94
95            if len(self._experts[bid]) >= n_experts * 3:
96                # merge the experts into a single 3d tensor
97                for w_name in ["down_proj", "gate_proj", "up_proj"]:
98                    datas: list[Tensor] = []
99
100                    for xid in range(n_experts):
101                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
102                        datas.append(self._experts[bid][ename])
103                        del self._experts[bid][ename]
104
105                    data_torch = torch.stack(datas, dim=0)
106
107                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
108
109                    yield from super().modify_tensors(data_torch, merged_name, bid)
110                return
111            else:
112                return
113
114        yield from super().modify_tensors(data_torch, name, bid)
115
116    # Copied from: Qwen2MoeModel
117    def prepare_tensors(self):
118        super().prepare_tensors()
119
120        if self._experts is not None:
121            # flatten `list[dict[str, Tensor]]` into `list[str]`
122            experts = [k for d in self._experts for k in d.keys()]
123            if len(experts) > 0:
124                raise ValueError(f"Unprocessed experts: {experts}")