Felipe97/llama-cpp-compiled

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 12from .llama import LlamaModel13 14 15@ModelBase.register("OlmoForCausalLM")16@ModelBase.register("OLMoForCausalLM")17@ModelBase.example("allenai/OLMo-1.7-7B-hf")18class OlmoModel(TextModel):19    model_arch = gguf.MODEL_ARCH.OLMO20 21    def set_gguf_parameters(self):22        super().set_gguf_parameters()23        self.gguf_writer.add_layer_norm_eps(1e-5)24        clip_qkv = self.hparams.get("clip_qkv")25        if clip_qkv is not None:26            self.gguf_writer.add_clamp_kqv(clip_qkv)27 28    # Same as super class, but permuting q_proj, k_proj29    # Copied from: LlamaModel30    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:31        n_head = self.hparams["num_attention_heads"]32        n_kv_head = self.hparams.get("num_key_value_heads")33 34        if name.endswith("q_proj.weight"):35            data_torch = LlamaModel.permute(data_torch, n_head, n_head)36        if name.endswith("k_proj.weight"):37            data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)38 39        yield from super().modify_tensors(data_torch, name, bid)40 41 42@ModelBase.register("SeedOssForCausalLM")43@ModelBase.example("ByteDance-Seed/Seed-OSS-36B-Instruct")44class SeedOssModel(TextModel):45    model_arch = gguf.MODEL_ARCH.SEED_OSS46 47 48@ModelBase.register("Olmo2ForCausalLM")49@ModelBase.register("Olmo3ForCausalLM")50@ModelBase.example("allenai/OLMo-2-1124-7B-Instruct", "allenai/Olmo-3-7B-Instruct")51class Olmo2Model(TextModel):52    model_arch = gguf.MODEL_ARCH.OLMO253 54    def set_gguf_parameters(self):55        super().set_gguf_parameters()56 57        if "sliding_window" in self.hparams:58            self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])59 60            sliding_window_pattern = []61            if "layer_types" in self.hparams:62                sliding_window_pattern = [t == "sliding_attention" for t in self.hparams["layer_types"]]63            else:64                # Olmo2 does not use sliding window attention.65                # Olmo3 defaults to using sliding window for all layers except every 4th.66                for i in range(self.hparams["num_hidden_layers"]):67                    sliding_window_pattern.append((i + 1) % 4 != 0)68 69            self.gguf_writer.add_sliding_window_pattern(sliding_window_pattern)70 71 72@ModelBase.register("OlmoeForCausalLM")73@ModelBase.example("allenai/OLMoE-1B-7B-0924")74class OlmoeModel(TextModel):75    model_arch = gguf.MODEL_ARCH.OLMOE76 77    def set_gguf_parameters(self):78        super().set_gguf_parameters()79        self.gguf_writer.add_layer_norm_rms_eps(1e-5)80 81    _experts: list[dict[str, Tensor]] | None = None82 83    # Copied from: Qwen2MoeModel84    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:85        # process the experts separately86        if name.find("experts") != -1:87            n_experts = self.find_hparam(["num_local_experts", "num_experts"])88            assert bid is not None89 90            if self._experts is None:91                self._experts = [{} for _ in range(self.block_count)]92 93            self._experts[bid][name] = data_torch94 95            if len(self._experts[bid]) >= n_experts * 3:96                # merge the experts into a single 3d tensor97                for w_name in ["down_proj", "gate_proj", "up_proj"]:98                    datas: list[Tensor] = []99 100                    for xid in range(n_experts):101                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"102                        datas.append(self._experts[bid][ename])103                        del self._experts[bid][ename]104 105                    data_torch = torch.stack(datas, dim=0)106 107                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"108 109                    yield from super().modify_tensors(data_torch, merged_name, bid)110                return111            else:112                return113 114        yield from super().modify_tensors(data_torch, name, bid)115 116    # Copied from: Qwen2MoeModel117    def prepare_tensors(self):118        super().prepare_tensors()119 120        if self._experts is not None:121            # flatten `list[dict[str, Tensor]]` into `list[str]`122            experts = [k for d in self._experts for k in d.keys()]123            if len(experts) > 0:124                raise ValueError(f"Unprocessed experts: {experts}")125