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 →
1.8 KiB50 linesraw
1from __future__ import annotations
2
3from typing import Iterable, TYPE_CHECKING
4
5if TYPE_CHECKING:
6    from torch import Tensor
7
8from .base import ModelBase, TextModel, gguf
9
10
11@ModelBase.register("BitnetForCausalLM", "BitNetForCausalLM")
12@ModelBase.example("microsoft/bitnet-b1.58-2B-4T")
13class BitnetModel(TextModel):
14    model_arch = gguf.MODEL_ARCH.BITNET
15
16    def set_vocab(self):
17        self._set_vocab_sentencepiece()
18
19    def set_gguf_parameters(self):
20        super().set_gguf_parameters()
21        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR)
22        self.gguf_writer.add_rope_scaling_factor(1.0)
23
24    def weight_quant(self, weight: Tensor) -> Tensor:
25        dtype = weight.dtype
26        weight = weight.float()
27        scale = weight.abs().mean().clamp(min=1e-5)
28        iscale = 1 / scale
29        # TODO: multiply by the scale directly instead of inverting it twice
30        # (this is also unnecessarily doubly inverted upstream)
31        # ref: https://huggingface.co/1bitLLM/bitnet_b1_58-3B/blob/af89e318d78a70802061246bf037199d2fb97020/utils_quant.py#L10
32        result = (weight * iscale).round().clamp(-1, 1) / iscale
33        return result.type(dtype)
34
35    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
36        new_name = self.map_tensor_name(name)
37
38        if any(self.match_model_tensor_name(new_name, key, bid) for key in [
39            gguf.MODEL_TENSOR.ATTN_Q,
40            gguf.MODEL_TENSOR.ATTN_K,
41            gguf.MODEL_TENSOR.ATTN_V,
42            gguf.MODEL_TENSOR.ATTN_OUT,
43            gguf.MODEL_TENSOR.FFN_UP,
44            gguf.MODEL_TENSOR.FFN_DOWN,
45            gguf.MODEL_TENSOR.FFN_GATE,
46        ]):
47            # transform weight into 1/0/-1 (in fp32)
48            data_torch = self.weight_quant(data_torch)
49
50        yield from super().modify_tensors(data_torch, name, bid)