Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("LLaDAModelLM")14class LLaDAModel(TextModel):15    model_arch = gguf.MODEL_ARCH.LLADA16    undo_permute = True17 18    def get_vocab_base(self) -> tuple[list[str], list[int], str]:19        tokens: list[str] = []20        toktypes: list[int] = []21 22        from transformers import AutoTokenizer23        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)24 25        vocab_dict = tokenizer.get_vocab()  # ty: ignore[unresolved-attribute]26        vocab_size = self.hparams.get("vocab_size", len(vocab_dict))27        assert max(vocab_dict.values()) < vocab_size28 29        tokpre = self.get_vocab_base_pre(tokenizer)30 31        reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab_dict.items()}32        added_vocab = tokenizer.get_added_vocab()  # ty: ignore[unresolved-attribute]33 34        for i in range(vocab_size):35            if i not in reverse_vocab:36                tokens.append(f"[PAD{i}]")37                toktypes.append(gguf.TokenType.UNUSED)38            elif reverse_vocab[i] in added_vocab:39                tokens.append(reverse_vocab[i])40                # Check if it's a special token - treat special tokens as CONTROL tokens41                if hasattr(tokenizer, 'added_tokens_decoder') and i in tokenizer.added_tokens_decoder:42                    if tokenizer.added_tokens_decoder[i].special:43                        toktypes.append(gguf.TokenType.CONTROL)44                    else:45                        toktypes.append(gguf.TokenType.USER_DEFINED)46                else:47                    # Fallback: treat all added vocab as control tokens for special tokens like <|im_start|>48                    toktypes.append(gguf.TokenType.CONTROL)49            else:50                tokens.append(reverse_vocab[i])51                toktypes.append(gguf.TokenType.NORMAL)52 53        return tokens, toktypes, tokpre54 55    def set_vocab(self):56        self._set_vocab_gpt2()57 58        # LLaDA specific parameters59        self.gguf_writer.add_add_bos_token(True)60 61    def set_gguf_parameters(self):62        super().set_gguf_parameters()63        self._try_set_pooling_type()64 65        # Add parameters similar to LlamaModel66        hparams = self.hparams67        self.gguf_writer.add_vocab_size(hparams["vocab_size"])68 69        if (rope_dim := hparams.get("head_dim")) is None:70            n_heads = hparams.get("num_attention_heads", hparams.get("n_heads"))71            assert n_heads is not None72            rope_dim = hparams.get("hidden_size", hparams.get("d_model")) // n_heads73        self.gguf_writer.add_rope_dimension_count(rope_dim)74 75        # Set context length for LLaDA76        context_length = self.hparams.get("max_sequence_length", 4096)77        self.gguf_writer.add_context_length(context_length)78 79        # Set embedding length (dimension size)80        embedding_length = self.hparams.get("d_model", 4096)81        self.gguf_writer.add_embedding_length(embedding_length)82 83        # Set feed forward length (MLP hidden size)84        feed_forward_length = self.hparams.get("mlp_hidden_size", 12288)85        self.gguf_writer.add_feed_forward_length(feed_forward_length)86 87        # LLaDA models use non-causal attention for diffusion, similar to Dream88        self.gguf_writer.add_causal_attention(False)89 90        # LLaDA models don't shift their logits91        self.gguf_writer.add_diffusion_shift_logits(False)92 93    @staticmethod94    def permute(weights: Tensor, n_head: int, n_head_kv: int | None):95        if n_head_kv is not None and n_head != n_head_kv:96            n_head = n_head_kv97        return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])98                .swapaxes(1, 2)99                .reshape(weights.shape))100 101    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:102        n_head = self.hparams.get("num_attention_heads", self.hparams.get("n_heads"))103        assert n_head is not None104        n_kv_head = self.hparams.get("num_key_value_heads", self.hparams.get("n_kv_heads"))105 106        if self.undo_permute:107            if name.endswith(("q_proj.weight", "q_proj.bias")):108                data_torch = LLaDAModel.permute(data_torch, n_head, n_head)109            if name.endswith(("k_proj.weight", "k_proj.bias")):110                data_torch = LLaDAModel.permute(data_torch, n_head, n_kv_head)111 112        # LLaDA model tensors should be mapped directly since it's the base model113        yield from super().modify_tensors(data_torch, name, bid)114 115 116@ModelBase.register("LLaDAMoEModel", "LLaDAMoEModelLM")117class LLaDAMoEModel(TextModel):118    model_arch = gguf.MODEL_ARCH.LLADA_MOE119 120    def set_gguf_parameters(self):121        super().set_gguf_parameters()122        if (expert_intermediate_size := self.hparams.get("expert_intermediate_size")) is not None:123            self.gguf_writer.add_expert_feed_forward_length(expert_intermediate_size)124 125        self.gguf_writer.add_mask_token_id(156895)126        self.gguf_writer.add_causal_attention(False)127        self.gguf_writer.add_diffusion_shift_logits(False)128 129    _experts: list[dict[str, Tensor]] | None = None130 131    # Copied from: Qwen2MoeModel132    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:133        # process the experts separately134        if name.find("experts") != -1:135            n_experts = self.find_hparam(["num_local_experts", "num_experts"])136            assert bid is not None137 138            if self._experts is None:139                self._experts = [{} for _ in range(self.block_count)]140 141            self._experts[bid][name] = data_torch142 143            if len(self._experts[bid]) >= n_experts * 3:144                # merge the experts into a single 3d tensor145                for w_name in ["down_proj", "gate_proj", "up_proj"]:146                    datas: list[Tensor] = []147 148                    for xid in range(n_experts):149                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"150                        datas.append(self._experts[bid][ename])151                        del self._experts[bid][ename]152 153                    data_torch = torch.stack(datas, dim=0)154 155                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"156 157                    yield from super().modify_tensors(data_torch, merged_name, bid)158                return159            else:160                return161 162        yield from super().modify_tensors(data_torch, name, bid)163 164    # Copied from: Qwen2MoeModel165    def prepare_tensors(self):166        super().prepare_tensors()167 168        if self._experts is not None:169            # flatten `list[dict[str, Tensor]]` into `list[str]`170            experts = [k for d in self._experts for k in d.keys()]171            if len(experts) > 0:172                raise ValueError(f"Unprocessed experts: {experts}")173