Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Callable, Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("Rwkv6ForCausalLM")14class Rwkv6Model(TextModel):15    model_arch = gguf.MODEL_ARCH.RWKV616 17    def set_vocab(self):18        self._set_vocab_rwkv_world()19 20    def set_gguf_parameters(self):21        head_size = self.hparams["head_size"]22        hidden_size = self.hparams["hidden_size"]23        layer_norm_eps = self.hparams["layer_norm_epsilon"]24        rescale_every_n_layers = self.hparams["rescale_every"]25        intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else int((hidden_size * 3.5) // 32 * 32)26        time_mix_extra_dim = 64 if hidden_size == 4096 else 3227        time_decay_extra_dim = 128 if hidden_size == 4096 else 6428 29        # RWKV isn't context limited30        self.gguf_writer.add_context_length(1048576)31        self.gguf_writer.add_embedding_length(hidden_size)32        self.gguf_writer.add_block_count(self.block_count)33        self.gguf_writer.add_layer_norm_eps(layer_norm_eps)34        self.gguf_writer.add_rescale_every_n_layers(rescale_every_n_layers)35        self.gguf_writer.add_wkv_head_size(head_size)36        self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim)37        self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim)38        self.gguf_writer.add_feed_forward_length(intermediate_size)39        self.gguf_writer.add_file_type(self.ftype)40 41        # required by llama.cpp, unused42        self.gguf_writer.add_head_count(0)43 44    lerp_weights: dict[int, dict[str, Tensor]] = {}45 46    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:47        new_name = self.map_tensor_name(name)48 49        if not (new_name.endswith(".weight") or new_name.endswith(".bias")):50            new_name += ".weight"51 52        if new_name.endswith("time_mix_w1.weight") or new_name.endswith("time_mix_decay_w1.weight") or new_name.endswith("time_mix_decay_w2.weight"):53            data_torch = data_torch.transpose(0, 1)54 55        if new_name.endswith("time_mix_w2.weight"):56            data_torch = data_torch.permute(0, 2, 1)57 58        if new_name.endswith("time_mix_decay.weight") or "lerp" in new_name:59            data_torch = data_torch.squeeze()60 61        try:62            rescale_every_n_layers = self.hparams["rescale_every"]63            if rescale_every_n_layers > 0:64                if new_name.endswith("time_mix_output.weight") or new_name.endswith("channel_mix_value.weight"):65                    data_torch = data_torch.div_(2 ** int(bid // rescale_every_n_layers))66        except KeyError:67            pass68 69        # concat time_mix_lerp weights to reduce some cpu overhead70        # also reduces the number of tensors in the model71        if bid is not None and "time_mix_lerp" in new_name and "time_mix_lerp_x" not in new_name:72            try:73                self.lerp_weights[bid][new_name] = data_torch74            except KeyError:75                self.lerp_weights[bid] = {new_name: data_torch}76            if all(f"blk.{bid}.time_mix_lerp_{i}.weight" in self.lerp_weights[bid].keys() for i in ["w", "k", "v", "r", "g"]):77                new_name = f"blk.{bid}.time_mix_lerp_fused.weight"78                data = torch.stack([self.lerp_weights[bid][f"blk.{bid}.time_mix_lerp_{i}.weight"].unsqueeze(0) for i in ["w", "k", "v", "r", "g"]], dim=0).unsqueeze(1)79                yield (new_name, data)80            return81 82        yield (new_name, data_torch)83 84 85@ModelBase.register("RWKV6Qwen2ForCausalLM")86class RWKV6Qwen2Model(Rwkv6Model):87    model_arch = gguf.MODEL_ARCH.RWKV6QWEN288 89    def set_vocab(self):90        try:91            self._set_vocab_sentencepiece()92        except FileNotFoundError:93            self._set_vocab_gpt2()94 95    def set_gguf_parameters(self):96        num_attention_heads = self.hparams["num_attention_heads"]97        num_key_value_heads = self.hparams["num_key_value_heads"]98        hidden_size = self.hparams["hidden_size"]99        head_size = hidden_size // num_attention_heads100        rms_norm_eps = self.hparams["rms_norm_eps"]101        intermediate_size = self.hparams["intermediate_size"]102        time_mix_extra_dim = self.hparams.get("lora_rank_tokenshift", 64 if hidden_size >= 4096 else 32)103        time_decay_extra_dim = self.hparams.get("lora_rank_decay", 128 if hidden_size >= 4096 else 64)104 105        # RWKV isn't context limited106        self.gguf_writer.add_context_length(1048576)107        self.gguf_writer.add_embedding_length(hidden_size)108        self.gguf_writer.add_block_count(self.block_count)109        self.gguf_writer.add_wkv_head_size(head_size)110        self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim)111        self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim)112        self.gguf_writer.add_feed_forward_length(intermediate_size)113        self.gguf_writer.add_file_type(self.ftype)114 115        # special parameters for time_mixing in RWKV6QWEN2116        self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)117        self.gguf_writer.add_token_shift_count(1)118        # RWKV6QWEN2 use grouped key/value like GQA119        self.gguf_writer.add_head_count_kv(num_key_value_heads)120 121        # required by llama.cpp, unused122        self.gguf_writer.add_head_count(0)123 124    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:125        for new_name, data in super().modify_tensors(data_torch, name, bid):126            if "time_mix_w1" in new_name or "time_mix_w2" in new_name:127                data = data.view(5, -1, data.shape[-1])128                # rwkv6qwen2 has a different order of rkvwg instead of the original wkvrg129                # permute them here to avoid code changes130                data = torch.stack([data[3], data[1], data[2], data[0], data[4]], dim=0).view(-1, data.shape[-1])131                if "w2" in new_name:132                    data = data.view(5, -1, data.shape[-1])133                yield (new_name, data)134                continue135            yield (new_name, data)136 137 138@ModelBase.register("Rwkv7ForCausalLM", "RWKV7ForCausalLM")139class Rwkv7Model(TextModel):140    model_arch = gguf.MODEL_ARCH.RWKV7141 142    def set_vocab(self):143        self._set_vocab_rwkv_world()144 145    def calc_lora_rank(self, hidden_size, exponent, multiplier):146        return max(1, round(hidden_size ** exponent * multiplier / 32)) * 32147 148    def set_gguf_parameters(self):149        try:150            head_size = self.hparams["head_size"]151            layer_norm_eps = self.hparams["layer_norm_epsilon"]152        except KeyError:153            head_size = self.hparams["head_dim"]154            layer_norm_eps = self.hparams["norm_eps"]155        hidden_size = self.hparams["hidden_size"]156        intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else (hidden_size * 4)157 158        # ICLR: In-Context-Learning-Rate159        try:160            lora_rank_decay = self.hparams["lora_rank_decay"] if self.hparams["lora_rank_decay"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)161            lora_rank_iclr = self.hparams["lora_rank_iclr"] if self.hparams["lora_rank_iclr"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)162            lora_rank_value_residual_mix = self.hparams["lora_rank_value_residual_mix"] if self.hparams["lora_rank_value_residual_mix"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3)163            lora_rank_gate = self.hparams["lora_rank_gate"] if self.hparams["lora_rank_gate"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6)164        except KeyError:165            lora_rank_decay = self.hparams["decay_low_rank_dim"] if self.hparams["decay_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)166            lora_rank_iclr = self.hparams["a_low_rank_dim"] if self.hparams["a_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)167            lora_rank_value_residual_mix = self.hparams["v_low_rank_dim"] if self.hparams["v_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3)168            lora_rank_gate = self.hparams["gate_low_rank_dim"] if self.hparams["gate_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6)169 170        # RWKV isn't context limited171        self.gguf_writer.add_context_length(1048576)172        self.gguf_writer.add_embedding_length(hidden_size)173        self.gguf_writer.add_block_count(self.block_count)174        self.gguf_writer.add_layer_norm_eps(layer_norm_eps)175        self.gguf_writer.add_wkv_head_size(head_size)176        self.gguf_writer.add_decay_lora_rank(lora_rank_decay)177        self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr)178        self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix)179        self.gguf_writer.add_gate_lora_rank(lora_rank_gate)180        self.gguf_writer.add_feed_forward_length(intermediate_size)181        self.gguf_writer.add_file_type(self.ftype)182 183        # required by llama.cpp, unused184        self.gguf_writer.add_head_count(0)185 186    lerp_weights: dict[int, dict[str, Tensor]] = {}187    lora_needs_transpose: bool = True188 189    @classmethod190    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:191        name, gen = item192 193        # unify tensor names here to make life easier194        name = name.replace("blocks", "layers").replace("ffn", "feed_forward")195        name = name.replace("self_attn", "attention").replace("attn", "attention")196        name = name.replace("time_mixer.", "")197 198        name = name.replace("feed_forward_norm", "ln2")199        name = name.replace("g_norm", "ln_x")200 201        return super().filter_tensors((name, gen))202 203    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:204        # lora layer names in fla-hub's impl205        if "_lora.lora" in name:206            self.lora_needs_transpose = False207        name = name.replace("_lora.lora.0.weight", "1.weight")208        name = name.replace("_lora.lora.2.weight", "2.weight")209        name = name.replace("_lora.lora.2.bias", "0.weight")210 211        if "attention.v" in name and "value" not in self.map_tensor_name(name) and bid == 0:212            # some models have dummy v0/v1/v2 on first layer while others don't213            # ignore them all since they are not used214            return215 216        wkv_has_gate = self.hparams.get("wkv_has_gate", True)217        lerp_list = ["r", "w", "k", "v", "a", "g"] if wkv_has_gate else ["r", "w", "k", "v", "a"]218 219        if bid is not None and "attention.x_" in name:220            if "attention.x_x" in name:221                # already concatenated222                new_name = f"blk.{bid}.time_mix_lerp_fused.weight"223                data = data_torch.reshape(len(lerp_list), 1, 1, -1)224                yield (new_name, data)225            else:226                try:227                    self.lerp_weights[bid][name] = data_torch228                except KeyError:229                    self.lerp_weights[bid] = {name: data_torch}230                if all(f"model.layers.{bid}.attention.x_{i}" in self.lerp_weights[bid].keys() for i in lerp_list):231                    new_name = f"blk.{bid}.time_mix_lerp_fused.weight"232                    data = torch.stack([self.lerp_weights[bid][f"model.layers.{bid}.attention.x_{i}"] for i in lerp_list], dim=0)233                    yield (new_name, data)234            return235        else:236            data_torch = data_torch.squeeze()237            new_name = self.map_tensor_name(name)238 239            if not (new_name.endswith(".weight") or new_name.endswith(".bias")):240                new_name += ".weight"241 242            if self.lora_needs_transpose and any(243                new_name.endswith(t) for t in [244                    "time_mix_w1.weight", "time_mix_w2.weight",245                    "time_mix_a1.weight", "time_mix_a2.weight",246                    "time_mix_v1.weight", "time_mix_v2.weight",247                    "time_mix_g1.weight", "time_mix_g2.weight",248                ]249            ):250                data_torch = data_torch.transpose(0, 1)251 252            if 'r_k' in new_name:253                data_torch = data_torch.flatten()254 255            if bid == 0 and "time_mix_a" in new_name:256                # dummy v0/v1/v2 on first layer257                # easiest way to make llama happy258                yield (new_name.replace("time_mix_a", "time_mix_v"), data_torch)259 260            yield (new_name, data_torch)261 262 263@ModelBase.register("RwkvHybridForCausalLM")264class ARwkv7Model(Rwkv7Model):265    model_arch = gguf.MODEL_ARCH.ARWKV7266 267    def set_vocab(self):268        try:269            self._set_vocab_sentencepiece()270        except FileNotFoundError:271            self._set_vocab_gpt2()272 273    def set_gguf_parameters(self):274        hidden_size = self.hparams["hidden_size"]275        head_size = self.hparams["head_size"]276        rms_norm_eps = self.hparams["rms_norm_eps"]277        intermediate_size = self.hparams["intermediate_size"]278        wkv_has_gate = self.hparams["wkv_has_gate"]279        assert self.hparams["wkv_version"] == 7280 281        # ICLR: In-Context-Learning-Rate282        lora_rank_decay = 64283        lora_rank_iclr = 64284        lora_rank_value_residual_mix = 32285        lora_rank_gate = 128 if wkv_has_gate else 0286 287        # RWKV isn't context limited288        self.gguf_writer.add_context_length(1048576)289        self.gguf_writer.add_embedding_length(hidden_size)290        self.gguf_writer.add_block_count(self.block_count)291        self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)292        self.gguf_writer.add_wkv_head_size(head_size)293        self.gguf_writer.add_decay_lora_rank(lora_rank_decay)294        self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr)295        self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix)296        self.gguf_writer.add_gate_lora_rank(lora_rank_gate)297        self.gguf_writer.add_feed_forward_length(intermediate_size)298        self.gguf_writer.add_file_type(self.ftype)299        self.gguf_writer.add_token_shift_count(1)300 301        # required by llama.cpp, unused302        self.gguf_writer.add_head_count(0)303