Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf9 10 11@ModelBase.register("BitnetForCausalLM", "BitNetForCausalLM")12@ModelBase.example("microsoft/bitnet-b1.58-2B-4T")13class BitnetModel(TextModel):14    model_arch = gguf.MODEL_ARCH.BITNET15 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.dtype26        weight = weight.float()27        scale = weight.abs().mean().clamp(min=1e-5)28        iscale = 1 / scale29        # TODO: multiply by the scale directly instead of inverting it twice30        # (this is also unnecessarily doubly inverted upstream)31        # ref: https://huggingface.co/1bitLLM/bitnet_b1_58-3B/blob/af89e318d78a70802061246bf037199d2fb97020/utils_quant.py#L1032        result = (weight * iscale).round().clamp(-1, 1) / iscale33        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)51