88 lines
3.5 KiB
Python
88 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Iterable, TYPE_CHECKING, cast
|
|
|
|
import torch
|
|
|
|
if TYPE_CHECKING:
|
|
from torch import Tensor
|
|
|
|
from .base import LazyTorchTensor, ModelBase, TextModel, gguf
|
|
|
|
|
|
@ModelBase.register("MapleForCausalLM")
|
|
@ModelBase.example("deepgrove/maple-preview")
|
|
class MapleModel(TextModel):
|
|
model_arch = gguf.MODEL_ARCH.MAPLE
|
|
|
|
def set_gguf_parameters(self):
|
|
super().set_gguf_parameters()
|
|
hparams = self.hparams
|
|
|
|
assert hparams["hidden_act"] == "silu"
|
|
assert hparams.get("num_shared_experts", 0) == 0
|
|
assert hparams.get("norm_topk_prob", True)
|
|
assert hparams.get("nope_on_global_attention", False)
|
|
|
|
head_dim = hparams.get("head_dim", hparams["hidden_size"] // hparams["num_attention_heads"])
|
|
partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 1.0)
|
|
|
|
self.gguf_writer.add_vocab_size(hparams["vocab_size"])
|
|
self.gguf_writer.add_rope_dimension_count(int(head_dim * partial_rotary_factor))
|
|
self.gguf_writer.add_sliding_window(hparams["sliding_window"])
|
|
self.gguf_writer.add_sliding_window_pattern([layer_type == "sliding_attention" for layer_type in hparams["layer_types"]])
|
|
self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
|
|
# the reference clamps the MoE SwiGLU gate/up at 7.0 (modeling_maple.py)
|
|
self.gguf_writer.add_swiglu_clamp_exp([7.0] * self.block_count)
|
|
|
|
_experts: list[dict[str, Tensor]] | None = None
|
|
|
|
@staticmethod
|
|
def _stack_experts(tensors: list[Tensor]) -> Tensor:
|
|
shape = (len(tensors), *tensors[0].shape)
|
|
dtype = tensors[0].dtype
|
|
meta = LazyTorchTensor.meta_with_dtype_and_shape(dtype, shape)
|
|
|
|
# tensors goes through args, not the closure, so that `func` matches
|
|
# LazyBase's single-argument shape
|
|
def stack(ts: list[Tensor]) -> Tensor:
|
|
result = torch.empty(shape, dtype=dtype)
|
|
for expert_id, tensor in enumerate(ts):
|
|
result[expert_id].copy_(LazyTorchTensor.to_eager(tensor))
|
|
ts.clear()
|
|
return result
|
|
|
|
return cast(torch.Tensor, LazyTorchTensor(meta=meta, args=(tensors,), func=stack))
|
|
|
|
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
|
if "mlp.experts" in name:
|
|
n_experts = self.hparams["num_experts"]
|
|
assert bid is not None
|
|
|
|
if self._experts is None:
|
|
self._experts = [{} for _ in range(self.block_count)]
|
|
|
|
self._experts[bid][name] = data_torch
|
|
|
|
if len(self._experts[bid]) >= n_experts * 3:
|
|
for weight_name in ("down_proj", "gate_proj", "up_proj"):
|
|
tensors = []
|
|
|
|
for expert_id in range(n_experts):
|
|
expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight"
|
|
tensors.append(self._experts[bid].pop(expert_name))
|
|
|
|
merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight"
|
|
yield from super().modify_tensors(self._stack_experts(tensors), merged_name, bid)
|
|
return
|
|
|
|
yield from super().modify_tensors(data_torch, name, bid)
|
|
|
|
def prepare_tensors(self):
|
|
super().prepare_tensors()
|
|
|
|
if self._experts is not None:
|
|
experts = [name for layer in self._experts for name in layer]
|
|
if experts:
|
|
raise ValueError(f"Unprocessed experts: {experts}")
|