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
2.1 KiB54 linesraw
1from __future__ import annotations
2
3from typing import Iterable, TYPE_CHECKING
4
5import torch
6
7if TYPE_CHECKING:
8    from torch import Tensor
9
10from .base import LazyTorchTensor, ModelBase, TextModel, gguf
11
12
13@ModelBase.register("TalkieForCausalLM")
14@ModelBase.example("lewtun/talkie-1930-13b-it-hf")
15class TalkieModel(TextModel):
16    model_arch = gguf.MODEL_ARCH.TALKIE
17
18    def set_gguf_parameters(self):
19        super().set_gguf_parameters()
20        # Talkie used F.rms_norm without an explicit eps
21        self.gguf_writer.add_layer_norm_rms_eps(torch.finfo(torch.float32).eps)
22
23    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
24        prefix = f"model.blocks.{bid}." if bid is not None else ""
25        suffix = name.removeprefix(prefix)
26
27        if suffix == "attn_gain.a_g":
28            yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_OUT, bid, ".scale"), data_torch
29            return
30        elif suffix == "mlp_gain.a_g":
31            yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN, bid, ".scale"), data_torch
32            return
33        elif suffix == "lm_head_gain.w_g":
34            self.gguf_writer.add_logit_scale(LazyTorchTensor.to_eager(data_torch).item())
35            return
36        elif suffix in ("attn.attn_query.weight", "attn.attn_key.weight"):
37            # absorb inverse rope
38            head_dim = self.hparams["head_dim"]
39            shape = data_torch.shape
40            data_torch = torch.reshape(data_torch, (-1, head_dim, shape[-1]))
41            signs = torch.ones((1, head_dim, 1), dtype=data_torch.dtype)
42            signs[:, head_dim // 2 :, :] = -1
43            if self.lazy:
44                signs = LazyTorchTensor.from_eager(signs)
45            # (n_head, head_dim, n_in) -> (n_out, n_in)
46            data_torch = torch.reshape(data_torch * signs, shape)
47        elif suffix == "attn.head_gain.head_g":
48            # allow head gain to broadcast
49            data_torch = data_torch.unsqueeze(-1)
50
51        if not name.endswith(".weight"):
52            name += ".weight"
53
54        yield from super().modify_tensors(data_torch, name, bid)