Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING, cast4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import LazyTorchTensor, ModelBase, TextModel, gguf11 12 13@ModelBase.register("MapleForCausalLM")14@ModelBase.example("deepgrove/maple-preview")15class MapleModel(TextModel):16    model_arch = gguf.MODEL_ARCH.MAPLE17 18    def set_gguf_parameters(self):19        super().set_gguf_parameters()20        hparams = self.hparams21 22        assert hparams["hidden_act"] == "silu"23        assert hparams.get("num_shared_experts", 0) == 024        assert hparams.get("norm_topk_prob", True)25        assert hparams.get("nope_on_global_attention", False)26 27        head_dim = hparams.get("head_dim", hparams["hidden_size"] // hparams["num_attention_heads"])28        partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 1.0)29 30        self.gguf_writer.add_vocab_size(hparams["vocab_size"])31        self.gguf_writer.add_rope_dimension_count(int(head_dim * partial_rotary_factor))32        self.gguf_writer.add_sliding_window(hparams["sliding_window"])33        self.gguf_writer.add_sliding_window_pattern([layer_type == "sliding_attention" for layer_type in hparams["layer_types"]])34        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])35        # the reference clamps the MoE SwiGLU gate/up at 7.0 (modeling_maple.py)36        self.gguf_writer.add_swiglu_clamp_exp([7.0] * self.block_count)37 38    _experts: list[dict[str, Tensor]] | None = None39 40    @staticmethod41    def _stack_experts(tensors: list[Tensor]) -> Tensor:42        shape = (len(tensors), *tensors[0].shape)43        dtype = tensors[0].dtype44        meta = LazyTorchTensor.meta_with_dtype_and_shape(dtype, shape)45 46        # tensors goes through args, not the closure, so that `func` matches47        # LazyBase's single-argument shape48        def stack(ts: list[Tensor]) -> Tensor:49            result = torch.empty(shape, dtype=dtype)50            for expert_id, tensor in enumerate(ts):51                result[expert_id].copy_(LazyTorchTensor.to_eager(tensor))52            ts.clear()53            return result54 55        return cast(torch.Tensor, LazyTorchTensor(meta=meta, args=(tensors,), func=stack))56 57    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:58        if "mlp.experts" in name:59            n_experts = self.hparams["num_experts"]60            assert bid is not None61 62            if self._experts is None:63                self._experts = [{} for _ in range(self.block_count)]64 65            self._experts[bid][name] = data_torch66 67            if len(self._experts[bid]) >= n_experts * 3:68                for weight_name in ("down_proj", "gate_proj", "up_proj"):69                    tensors = []70 71                    for expert_id in range(n_experts):72                        expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight"73                        tensors.append(self._experts[bid].pop(expert_name))74 75                    merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight"76                    yield from super().modify_tensors(self._stack_experts(tensors), merged_name, bid)77            return78 79        yield from super().modify_tensors(data_torch, name, bid)80 81    def prepare_tensors(self):82        super().prepare_tensors()83 84        if self._experts is not None:85            experts = [name for layer in self._experts for name in layer]86            if experts:87                raise ValueError(f"Unprocessed experts: {experts}")88