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 →
5.4 KiB119 linesraw
1from __future__ import annotations
2
3from typing import Any, Iterable, TYPE_CHECKING
4
5if TYPE_CHECKING:
6    from torch import Tensor
7
8from .base import ModelBase, gguf
9
10from .llama import LlamaModel
11from .mamba import Mamba2Model
12
13
14@ModelBase.register("FalconH1ForCausalLM")
15@ModelBase.example("tiiuae/Falcon-H1-0.5B-Base")
16class FalconH1Model(Mamba2Model):
17    model_arch = gguf.MODEL_ARCH.FALCON_H1
18
19    def __init__(self, *args, **kwargs):
20        # Set the hparam prefixes for Falcon Mamba2
21        self.hparam_prefixes = ["mamba"]
22
23        # Initialize the base Mamba2Model
24        super().__init__(*args, **kwargs)
25
26        # Use Llama conversion for attention
27        self._transformer_model_class = LlamaModel
28
29        # n_group and d_inner are used during reshape_tensors for mamba2
30        self.n_group = self.find_hparam(["n_groups"])
31        self.d_inner = self.find_hparam(["mamba_d_ssm"])
32        self.d_head = self.find_hparam(["d_head"])
33
34        # Initialize any Falcon Mamba2 specific attributes
35        self.has_attention = True  # Falcon Mamba2 has attention components
36
37        # Load Falcon-H1 multipliers from hyperparameters
38        self.attention_in_multiplier = self.find_hparam(["attention_in_multiplier"], optional=True)
39        self.attention_out_multiplier = self.find_hparam(["attention_out_multiplier"], optional=True)
40        self.ssm_in_multiplier = self.find_hparam(["ssm_in_multiplier"], optional=True)
41        self.ssm_out_multiplier = self.find_hparam(["ssm_out_multiplier"], optional=True)
42        self.mlp_multipliers = self.find_hparam(["mlp_multipliers"], optional=True)
43        self.ssm_multipliers = self.find_hparam(["ssm_multipliers"], optional=True)
44        self.intermediate_size = self.find_hparam(["intermediate_size"])
45        self.key_multiplier = self.find_hparam(["key_multiplier"], optional=True)
46
47    def find_hparam(self, keys: Iterable[str], *args, **kwargs) -> Any:
48        prefixed = []
49        for pfx in self.hparam_prefixes:
50            prefixed.extend(
51                "_".join([pfx, k])
52                for k in keys
53            )
54        keys = list(keys) + prefixed
55        return super().find_hparam(keys, *args, **kwargs)
56
57    def set_vocab(self):
58        self._set_vocab_gpt2()
59
60    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
61        tensors = list(super().modify_tensors(data_torch, name, bid))
62        tensor = tensors[0][1]
63
64        if "down_proj" in name:
65            tensor = tensor  * self.mlp_multipliers[1]
66        elif "gate_proj" in name:
67            tensor = tensor * self.mlp_multipliers[0]
68        elif "k_proj" in name:
69            tensor = tensor * self.key_multiplier * self.attention_in_multiplier
70        elif "q_proj" in name:
71            tensor = tensor * self.attention_in_multiplier
72        elif "v_proj" in name:
73            tensor = tensor * self.attention_in_multiplier
74        elif "o_proj" in name:
75            tensor = tensor * self.attention_out_multiplier
76        elif "out_proj" in name:
77            tensor = tensor * self.ssm_out_multiplier
78        elif "in_proj" in name:
79            tensor = tensor * self.ssm_in_multiplier
80            zxbcdt_multipliers = self.hparams["ssm_multipliers"]
81            intermediate_size = self.hparams["mamba_d_ssm"]
82            groups_time_state_size = self.hparams["mamba_n_groups"] * self.hparams["mamba_d_state"]
83            tensor[:intermediate_size, :] *= zxbcdt_multipliers[0]
84            tensor[intermediate_size:2 * intermediate_size, :] *= zxbcdt_multipliers[1]
85            tensor[2 * intermediate_size:2 * intermediate_size + groups_time_state_size, :] *= zxbcdt_multipliers[2]
86            tensor[2 * intermediate_size + groups_time_state_size:2 * intermediate_size + 2 * groups_time_state_size, :] *= zxbcdt_multipliers[3]
87            tensor[2 * intermediate_size + 2 * groups_time_state_size:, :] *= zxbcdt_multipliers[4]
88        elif "lm_head" in name:
89            tensor = tensor * self.hparams["lm_head_multiplier"]
90        elif "embed_tokens" in name:
91            tensor = tensor * self.hparams["embedding_multiplier"]
92        elif "mamba.norm" in name:
93            tensor = tensor.reshape(self.n_group, self.d_inner // self.n_group)
94
95        tensors = [(tensors[0][0], tensor)]
96        return tensors
97
98    def set_gguf_parameters(self):
99        super().set_gguf_parameters()
100
101        ## General Params ##
102        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])
103        # Override some Mamba2 defaults
104        self.gguf_writer.add_block_count(self.block_count)
105        self.gguf_writer.add_context_length(self.hparams.get("max_position_embeddings", 0))
106        self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"])
107
108        ## Attention params ##
109        self.gguf_writer.add_head_count(self.hparams["num_attention_heads"]) # Override value 0 from Mamba2
110        self.gguf_writer.add_head_count_kv(self.hparams["num_key_value_heads"])
111        self.gguf_writer.add_key_length(self.hparams["head_dim"])
112        self.gguf_writer.add_value_length(self.hparams["head_dim"])
113
114        ## Validation ##
115        assert self.hparams.get("hidden_act") in [None, "silu"], "Only SILU activation supported"
116        assert self.d_inner % self.d_head == 0, f"SSM inner size {self.d_inner} not a multiple of head dim {self.d_head}"
117
118        # Add any other Falcon Mamba2 specific configuration
119        self.gguf_writer.add_rope_freq_base(self.rope_parameters["rope_theta"])