Buckets:
| 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 | |
| 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 | |
| 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}") | |
Xet Storage Details
- Size:
- 3.59 kB
- Xet hash:
- 06b4d9963a423d0f62f602dc0e6392039c58772155b45cddb1b0d74a7fecb545
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.