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 →
5.8 KiB131 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, TextModel, gguf, logger
11
12
13@ModelBase.register("GptOssForCausalLM")
14@ModelBase.example("openai/gpt-oss-20b")
15class GptOssModel(TextModel):
16    model_arch = gguf.MODEL_ARCH.GPT_OSS
17
18    # TODO: remove once MXFP4 is supported more generally
19    def dequant_model(self):
20        if self._is_mxfp4:
21            return
22        return super().dequant_model()
23
24    def transform_nibble_layout(self, tensor):
25        assert tensor.dtype == torch.uint8
26        assert tensor.shape[-1] == 16
27        # swap nibbles
28        t_lo = tensor & 0x0F
29        t_hi = tensor & 0xF0
30        t_swapped = (t_lo << 4) | (t_hi >> 4)
31        tensor = t_swapped
32        # transform aaaa...bbbb... to abababab...
33        blk_a, blk_b = tensor.chunk(2, dim=-1)
34        # get a_
35        blk_a0 = (blk_a & 0xF0).view(-1, 1)
36        blk_a1 = (blk_a << 4).view(-1, 1)
37        blk_a = torch.stack((blk_a0, blk_a1), dim=2).view(tensor.shape)
38        # get _b
39        blk_b0 = (blk_b >> 4).view(-1, 1)
40        blk_b1 = (blk_b & 0x0F).view(-1, 1)
41        blk_b = torch.stack((blk_b0, blk_b1), dim=2).view(tensor.shape)
42        # swap once more
43        out = blk_a | blk_b
44        out_h = out & 0xF0
45        out_l = out & 0x0F
46        out = (out_h >> 4) | (out_l << 4)
47        return out
48
49    def repack_mxfp4(self, new_name: str, blocks: Tensor, scales: Tensor):
50        assert blocks.dtype == torch.uint8
51        assert scales.dtype == torch.uint8
52        scales = scales.unsqueeze(-1)
53        assert len(blocks.shape) == 4
54        assert len(scales.shape) == 4
55        blocks = self.transform_nibble_layout(blocks)
56        new_data = torch.concat((scales, blocks), dim=-1)
57        new_shape = [new_data.shape[0], new_data.shape[1], new_data.shape[2] * 32]
58        logger.info(f"Repacked {new_name} with shape {new_shape} and quantization MXFP4")
59        # flatten last dim
60        new_data = new_data.view(new_data.shape[0], new_data.shape[1], new_data.shape[2] * new_data.shape[3])
61        new_data = new_data.numpy()
62        self.gguf_writer.add_tensor(new_name, new_data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)
63
64    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
65        blocks0: Tensor = torch.zeros(1)
66        blocks1: Tensor = torch.zeros(1)
67        # we assume that tensors are loaded in the correct order
68        for name, data_torch in self.get_tensors():
69            if "mlp.experts.down_proj_blocks" in name:
70                blocks0 = data_torch
71            elif "mlp.experts.down_proj_scales" in name:
72                new_name = self.map_tensor_name(name.replace("_scales", ".weight"))
73                self.repack_mxfp4(new_name, blocks0, data_torch)
74            elif "mlp.experts.gate_up_proj_blocks" in name:
75                blocks0, blocks1 = data_torch[:, ::2, :, :], data_torch[:, 1::2, :, :]
76            elif "mlp.experts.gate_up_proj_scales" in name:
77                scales0, scales1 = data_torch[:, ::2, :], data_torch[:, 1::2, :]
78                new_name_gate = self.map_tensor_name(name.replace("gate_up_proj_scales", "gate_proj.weight"))
79                new_name_up = self.map_tensor_name(name.replace("gate_up_proj_scales", "up_proj.weight"))
80                self.repack_mxfp4(new_name_gate, blocks0, scales0)
81                self.repack_mxfp4(new_name_up, blocks1, scales1)
82        return []
83
84    @classmethod
85    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
86        name, gen = item
87
88        if "sinks" in name:
89            name += ".weight"
90
91        return super().filter_tensors((name, gen))
92
93    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
94        # correct naming for down_proj
95        if "down_proj" in name:
96            if name.endswith("_bias"):
97                name = name.replace("down_proj_bias", "down_proj.bias")
98            elif "_blocks" not in name and "_scales" not in name:
99                logger.warning(f"{name} is not in MXFP4, performance may be degraded")
100                name = name.replace("down_proj", "down_proj.weight")
101                data_torch = data_torch.transpose(-1, -2)
102            else:
103                # otherwise, it should already be repacked to ggml MXFP4 format
104                return
105
106        # split the gate_up into gate and up
107        if "gate_up_proj" in name:
108            if name.endswith("_bias"):
109                name_up = name.replace("gate_up_proj_bias", "up_proj.bias")
110                name_gate = name.replace("gate_up_proj_bias", "gate_proj.bias")
111                gate_proj_bias, up_proj_bias = data_torch[..., ::2], data_torch[..., 1::2]
112                yield from super().modify_tensors(gate_proj_bias, name_gate, bid)
113                yield from super().modify_tensors(up_proj_bias, name_up, bid)
114            elif "_blocks" not in name and "_scales" not in name:
115                logger.warning(f"{name} is not in MXFP4, performance may be degraded")
116                name_up = name.replace("gate_up_proj", "up_proj.weight")
117                name_gate = name.replace("gate_up_proj", "gate_proj.weight")
118                data_torch = data_torch.transpose(-1, -2)
119                gate_proj_weight, up_proj_weight = data_torch[:, ::2, :], data_torch[:, 1::2, :]
120                yield from super().modify_tensors(gate_proj_weight, name_gate, bid)
121                yield from super().modify_tensors(up_proj_weight, name_up, bid)
122        else:
123            yield from super().modify_tensors(data_torch, name, bid)
124
125    def set_vocab(self):
126        self._set_vocab_gpt2()
127
128    def set_gguf_parameters(self):
129        super().set_gguf_parameters()
130        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
131        self.gguf_writer.add_expert_feed_forward_length(self.hparams["intermediate_size"])