Brunobkr/llama.cpp_AlgMor24_github

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")12class BitnetModel(TextModel):13    model_arch = gguf.MODEL_ARCH.BITNET14 15    def set_vocab(self):16        self._set_vocab_sentencepiece()17 18    def set_gguf_parameters(self):19        super().set_gguf_parameters()20        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR)21        self.gguf_writer.add_rope_scaling_factor(1.0)22 23    def weight_quant(self, weight: Tensor) -> Tensor:24        dtype = weight.dtype25        weight = weight.float()26        scale = weight.abs().mean().clamp(min=1e-5)27        iscale = 1 / scale28        # TODO: multiply by the scale directly instead of inverting it twice29        # (this is also unnecessarily doubly inverted upstream)30        # ref: https://huggingface.co/1bitLLM/bitnet_b1_58-3B/blob/af89e318d78a70802061246bf037199d2fb97020/utils_quant.py#L1031        result = (weight * iscale).round().clamp(-1, 1) / iscale32        return result.type(dtype)33 34    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:35        new_name = self.map_tensor_name(name)36 37        if any(self.match_model_tensor_name(new_name, key, bid) for key in [38            gguf.MODEL_TENSOR.ATTN_Q,39            gguf.MODEL_TENSOR.ATTN_K,40            gguf.MODEL_TENSOR.ATTN_V,41            gguf.MODEL_TENSOR.ATTN_OUT,42            gguf.MODEL_TENSOR.FFN_UP,43            gguf.MODEL_TENSOR.FFN_DOWN,44            gguf.MODEL_TENSOR.FFN_GATE,45        ]):46            # transform weight into 1/0/-1 (in fp32)47            data_torch = self.weight_quant(data_torch)48 49        yield from super().modify_tensors(data_torch, name, bid)50