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 →
2.5 KiB60 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, logger
9
10
11@ModelBase.register("BaichuanForCausalLM", "BaiChuanForCausalLM")
12@ModelBase.example("baichuan-inc/Baichuan2-7B-Chat", "baichuan-inc/Baichuan-7B")
13class BaichuanModel(TextModel):
14    model_arch = gguf.MODEL_ARCH.BAICHUAN
15
16    def set_vocab(self):
17        self._set_vocab_sentencepiece()
18
19    def set_gguf_parameters(self):
20        super().set_gguf_parameters()
21
22        self.gguf_writer.add_tensor_data_layout("Meta AI original pth")
23        self.gguf_writer.add_rope_dimension_count(self.hparams["hidden_size"] // self.hparams["num_attention_heads"])
24
25    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
26        head_count = self.hparams["num_attention_heads"]
27        head_count_kv = self.hparams.get("num_key_value_heads", head_count)
28
29        if bid is not None and name == f"model.layers.{bid}.self_attn.W_pack.weight":
30            logger.info(f"Unpacking and permuting layer {bid}")
31            yield from [
32                (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid),
33                    self._reverse_hf_permute_part(data_torch, 0, head_count, head_count)),
34                (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid),
35                    self._reverse_hf_permute_part(data_torch, 1, head_count, head_count_kv)),
36                (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid),
37                    self._reverse_hf_part(data_torch, 2)),
38            ]
39        else:
40            yield from self.modify_tensors(data_torch, self.map_tensor_name(name), bid)
41
42    def _reverse_hf_permute(self, weights: Tensor, n_head: int, n_kv_head: int | None = None) -> Tensor:
43        if n_kv_head is not None and n_head != n_kv_head:
44            n_head //= n_kv_head
45
46        return (
47            weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])
48            .swapaxes(1, 2)
49            .reshape(weights.shape)
50        )
51
52    def _reverse_hf_permute_part(
53        self, weights: Tensor, n_part: int, n_head: int, n_head_kv: int | None = None,
54    ) -> Tensor:
55        r = weights.shape[0] // 3
56        return self._reverse_hf_permute(weights[r * n_part:r * n_part + r, ...], n_head, n_head_kv)
57
58    def _reverse_hf_part(self, weights: Tensor, n_part: int) -> Tensor:
59        r = weights.shape[0] // 3
60        return weights[r * n_part:r * n_part + r, ...]