Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import json4 5from typing import Iterable, TYPE_CHECKING6 7import torch8 9if TYPE_CHECKING:10    from torch import Tensor11 12from .base import ModelBase, TextModel, gguf13 14 15@ModelBase.register("PlamoForCausalLM")16class PlamoModel(TextModel):17    model_arch = gguf.MODEL_ARCH.PLAMO18 19    def set_vocab(self):20        self._set_vocab_sentencepiece()21 22    def set_gguf_parameters(self):23        hparams = self.hparams24 25        self.gguf_writer.add_context_length(4096)  # not in config.json26        self.gguf_writer.add_embedding_length(hparams["hidden_size"])27        self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"])28        self.gguf_writer.add_block_count(self.block_count)29        self.gguf_writer.add_head_count(hparams["num_attention_heads"])30        self.gguf_writer.add_head_count_kv(5)  # hparams["num_key_value_heads"]) is wrong31        self.gguf_writer.add_layer_norm_rms_eps(hparams["rms_norm_eps"])32        self.gguf_writer.add_file_type(self.ftype)33 34    def shuffle_attn_q_weight(self, data_torch):35        assert data_torch.size() == (5120, 5120)36        data_torch = data_torch.reshape(8, 5, 128, 5120)37        data_torch = torch.permute(data_torch, (1, 0, 2, 3))38        data_torch = torch.reshape(data_torch, (5120, 5120))39        return data_torch40 41    def shuffle_attn_output_weight(self, data_torch):42        assert data_torch.size() == (5120, 5120)43        data_torch = data_torch.reshape(5120, 8, 5, 128)44        data_torch = torch.permute(data_torch, (0, 2, 1, 3))45        data_torch = torch.reshape(data_torch, (5120, 5120))46        return data_torch47 48    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:49        new_name = self.map_tensor_name(name)50 51        # shuffle for broadcasting of gqa in ggml_mul_mat52        if new_name.endswith("attn_q.weight"):53            data_torch = self.shuffle_attn_q_weight(data_torch)54        elif new_name.endswith("attn_output.weight"):55            data_torch = self.shuffle_attn_output_weight(data_torch)56 57        yield from super().modify_tensors(data_torch, name, bid)58 59 60@ModelBase.register("Plamo2ForCausalLM", "PLaMo2ForCausalLM")61class Plamo2Model(TextModel):62    model_arch = gguf.MODEL_ARCH.PLAMO263 64    def set_vocab(self):65        self._set_vocab_plamo()66 67    def set_gguf_parameters(self):68        hparams = self.hparams69        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])70 71        # Which layers are Mamba layers72        # PLaMo 2 uses mamba_step to indicate the pattern (e.g., 2 means every other layer)73        # This logic matches modeling_plamo.py's is_mamba function74        mamba_step = hparams.get("mamba_step", 2)75        mamba_enabled = hparams.get("mamba_enabled", True)76        num_key_value_heads = []77        num_attention_heads = []78 79        if mamba_enabled:80            for i in range(self.block_count):81                if self.block_count <= (mamba_step // 2):82                    # use attention in last layer83                    is_mamba = (i != self.block_count - 1)84                else:85                    is_mamba = (i % mamba_step) != (mamba_step // 2)86                if is_mamba:87                    num_key_value_heads.append(0)88                    num_attention_heads.append(0)89                else:90                    num_key_value_heads.append(hparams.get("num_key_value_heads", 4))91                    num_attention_heads.append(hparams.get("num_attention_heads", 32))92 93        if num_key_value_heads and num_attention_heads:94            self.gguf_writer.add_head_count_kv(num_key_value_heads)95            self.gguf_writer.add_head_count(num_attention_heads)96 97        self.gguf_writer.add_context_length(hparams.get("max_position_embeddings", 2048))98        self.gguf_writer.add_embedding_length(hparams.get("hidden_size", 4096))99        self.gguf_writer.add_key_length(hparams.get("hidden_size_per_head", 128))100        self.gguf_writer.add_value_length(hparams.get("hidden_size_per_head", 128))101        self.gguf_writer.add_block_count(self.block_count)102        self.gguf_writer.add_layer_norm_rms_eps(hparams.get("rms_norm_eps", 1e-06))103        self.gguf_writer.add_rope_freq_base(self.rope_parameters.get("rope_theta", 10000))104 105        # Mamba parameters106        self.gguf_writer.add_ssm_state_size(hparams.get("mamba_d_state", 64))107        self.gguf_writer.add_ssm_conv_kernel(hparams.get("mamba_d_conv", 4))108        self.gguf_writer.add_ssm_time_step_rank(hparams.get("mamba_num_heads", 64))109        intermediate_size = hparams.get("mamba_num_heads", 64) * hparams.get("hidden_size_per_head", 128)110        self.gguf_writer.add_ssm_inner_size(intermediate_size)111        self.gguf_writer.add_ssm_group_count(0)112 113        # MLP feed forward parameters (for attention layers)114        self.gguf_writer.add_feed_forward_length(hparams.get("intermediate_size", 13312))115        self.gguf_writer.add_file_type(self.ftype)116 117    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:118        if name.endswith(".A_log"):119            data_torch = -torch.exp(data_torch)120        elif name.endswith(".dt_bias"):121            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"122        elif name.endswith(".dt_norm_weight"):123            name = name.rpartition(".dt_norm_weight")[0] + ".dt_norm.weight"124        elif name.endswith(".B_norm_weight"):125            name = name.rpartition(".B_norm_weight")[0] + ".B_norm.weight"126        elif name.endswith(".C_norm_weight"):127            name = name.rpartition(".C_norm_weight")[0] + ".C_norm.weight"128        elif name.endswith(".k_weight"):129            name = name.rpartition(".k_weight")[0] + ".k.weight"130        elif name.endswith(".q_weight"):131            name = name.rpartition(".q_weight")[0] + ".q.weight"132        elif name.endswith(".conv1d.weight"):133            data_torch = torch.squeeze(data_torch)  # remove (, 1, )134            assert data_torch.ndim == 2135        elif name.endswith(".pre_mixer_norm.weight"):136            data_torch += 1.0137        elif name.endswith(".post_mixer_norm.weight"):138            data_torch += 1.0 / 5139        elif name.endswith(".pre_mlp_norm.weight"):140            data_torch += 1.0141        elif name.endswith(".post_mlp_norm.weight"):142            data_torch += 1.0 / (5**1.5)143        elif name.endswith(".norm.weight"):144            data_torch += 1.0145 146        yield from super().modify_tensors(data_torch, name, bid)147 148 149@ModelBase.register("Plamo3ForCausalLM", "PLaMo3ForCausalLM")150class Plamo3Model(TextModel):151    model_arch = gguf.MODEL_ARCH.PLAMO3152 153    def set_vocab(self):154        self._set_vocab_plamo()155 156        tokenizer_config_path = self.dir_model / "tokenizer_config.json"157        tokenizer_config = {}158 159        if tokenizer_config_path.is_file():160            with open(tokenizer_config_path, encoding="utf-8") as f:161                tokenizer_config = json.load(f)162 163        chat_template = tokenizer_config.get("chat_template")164        chat_template_jinja = self.dir_model / "chat_template.jinja"165 166        if chat_template_jinja.is_file():167            with open(chat_template_jinja, encoding="utf-8") as f:168                chat_template = f.read()169 170        if chat_template:171            self.gguf_writer.add_chat_template(chat_template)172 173    def set_gguf_parameters(self):174        super().set_gguf_parameters()175        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])176        if (sliding_window := self.find_hparam(["window_size", "sliding_window"], optional=True)) is not None:177            self.gguf_writer.add_sliding_window(sliding_window)178            self.gguf_writer.add_sliding_window_pattern(self.hparams["sliding_window_pattern"])179 180    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:181 182        if name.endswith(".pre_mixer_norm.weight"):183            data_torch = data_torch + 1.0184        elif name.endswith(".post_mixer_norm.weight"):185            data_torch = data_torch + 1.0 / 5186        elif name.endswith(".pre_mlp_norm.weight"):187            data_torch = data_torch + 1.0188        elif name.endswith(".post_mlp_norm.weight"):189            data_torch = data_torch + 1.0 / (5**1.5)190        elif name.endswith((".mixer.q_norm.weight", ".mixer.k_norm.weight")):191            data_torch = data_torch + 1.0192        elif name.endswith(".norm.weight"):193            data_torch = data_torch + 1.0194 195        yield from super().modify_tensors(data_torch, name, bid)196