char-fork/llama.cpp

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

Alexllama: add Maple 20B-A1B ternary MoE architecture (CPU) (#27000)3d10bcd19

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.5 KiB87 linesraw
1from __future__ import annotations
2
3from typing import Iterable, TYPE_CHECKING, cast
4
5import torch
6
7if TYPE_CHECKING:
8    from torch import Tensor
9
10from .base import LazyTorchTensor, ModelBase, TextModel, gguf
11
12
13@ModelBase.register("MapleForCausalLM")
14@ModelBase.example("deepgrove/maple-preview")
15class MapleModel(TextModel):
16    model_arch = gguf.MODEL_ARCH.MAPLE
17
18    def set_gguf_parameters(self):
19        super().set_gguf_parameters()
20        hparams = self.hparams
21
22        assert hparams["hidden_act"] == "silu"
23        assert hparams.get("num_shared_experts", 0) == 0
24        assert hparams.get("norm_topk_prob", True)
25        assert hparams.get("nope_on_global_attention", False)
26
27        head_dim = hparams.get("head_dim", hparams["hidden_size"] // hparams["num_attention_heads"])
28        partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 1.0)
29
30        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
31        self.gguf_writer.add_rope_dimension_count(int(head_dim * partial_rotary_factor))
32        self.gguf_writer.add_sliding_window(hparams["sliding_window"])
33        self.gguf_writer.add_sliding_window_pattern([layer_type == "sliding_attention" for layer_type in hparams["layer_types"]])
34        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
35        # the reference clamps the MoE SwiGLU gate/up at 7.0 (modeling_maple.py)
36        self.gguf_writer.add_swiglu_clamp_exp([7.0] * self.block_count)
37
38    _experts: list[dict[str, Tensor]] | None = None
39
40    @staticmethod
41    def _stack_experts(tensors: list[Tensor]) -> Tensor:
42        shape = (len(tensors), *tensors[0].shape)
43        dtype = tensors[0].dtype
44        meta = LazyTorchTensor.meta_with_dtype_and_shape(dtype, shape)
45
46        # tensors goes through args, not the closure, so that `func` matches
47        # LazyBase's single-argument shape
48        def stack(ts: list[Tensor]) -> Tensor:
49            result = torch.empty(shape, dtype=dtype)
50            for expert_id, tensor in enumerate(ts):
51                result[expert_id].copy_(LazyTorchTensor.to_eager(tensor))
52            ts.clear()
53            return result
54
55        return cast(torch.Tensor, LazyTorchTensor(meta=meta, args=(tensors,), func=stack))
56
57    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
58        if "mlp.experts" in name:
59            n_experts = self.hparams["num_experts"]
60            assert bid is not None
61
62            if self._experts is None:
63                self._experts = [{} for _ in range(self.block_count)]
64
65            self._experts[bid][name] = data_torch
66
67            if len(self._experts[bid]) >= n_experts * 3:
68                for weight_name in ("down_proj", "gate_proj", "up_proj"):
69                    tensors = []
70
71                    for expert_id in range(n_experts):
72                        expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight"
73                        tensors.append(self._experts[bid].pop(expert_name))
74
75                    merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight"
76                    yield from super().modify_tensors(self._stack_experts(tensors), merged_name, bid)
77            return
78
79        yield from super().modify_tensors(data_torch, name, bid)
80
81    def prepare_tensors(self):
82        super().prepare_tensors()
83
84        if self._experts is not None:
85            experts = [name for layer in self._experts for name in layer]
86            if experts:
87                raise ValueError(f"Unprocessed experts: {experts}")