Buckets:
| from __future__ import annotations | |
| from typing import Any, Iterable, TYPE_CHECKING | |
| import torch | |
| if TYPE_CHECKING: | |
| from torch import Tensor | |
| from .base import ModelBase, MmprojModel, TextModel, gguf | |
| from .muse_glimmer import MuseGlimmerVisionModel | |
| # The reference clamps the squared stream coefficient with this floor. | |
| RESIDUAL_GATE_EPS = 1e-3 | |
| # Block norms hold a delta around one; the output norm holds an absolute gain. | |
| OFFSET_NORM_SUFFIXES = ( | |
| "input_layernorm.weight", | |
| "pre_feedforward_layernorm.weight", | |
| "post_feedforward_layernorm.weight", | |
| "pre_expert_norm.weight", | |
| "post_expert_norm.weight", | |
| ) | |
| def _unpermute_for_rope(tensor: "Tensor", n_heads: int) -> "Tensor": | |
| """Invert transformers' `_permute_for_rope`: HF stores Q/K in rotate_half layout, | |
| llama.cpp consumes the interleaved (NORM) layout. Vision Q/K carry a bias too.""" | |
| if tensor.ndim == 1: | |
| dim1, = tensor.shape | |
| return tensor.view(n_heads, 2, dim1 // n_heads // 2).transpose(1, 2).reshape(dim1) | |
| dim1, dim2 = tensor.shape | |
| return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) | |
| class MuseSparkModel(TextModel): | |
| model_arch = gguf.MODEL_ARCH.MUSE_SPARK | |
| def set_vocab(self): | |
| self._set_vocab_gpt2() | |
| from transformers import AutoTokenizer | |
| tok = AutoTokenizer.from_pretrained(self.dir_model) | |
| eot_id = tok.convert_tokens_to_ids("<|eot|>") # ty: ignore[unresolved-attribute] | |
| if isinstance(eot_id, int) and eot_id >= 0: | |
| self.gguf_writer.add_eot_token_id(eot_id) | |
| def set_gguf_parameters(self): | |
| super().set_gguf_parameters() | |
| hparams = self.hparams | |
| # MetaP resolves to two scalars | |
| # The output multiplier lands here and the attention scale rides on the synthesized QK-norm weights | |
| self.gguf_writer.add_logit_scale(hparams["output_multiplier"]) | |
| self.gguf_writer.add_final_logit_softcapping(hparams["final_logit_softcapping"]) | |
| self.gguf_writer.add_post_norm_eps(hparams["post_norm_eps"]) | |
| # Only the sliding layers carry rope, so the graph takes it from the pattern | |
| sliding = [t == "sliding_attention" for t in hparams["layer_types"]] | |
| theta = hparams["rope_parameters"]["rope_theta"] | |
| assert sliding == [bool(t) for t in hparams["layer_rope_theta"]] | |
| assert all(t in (0, theta) for t in hparams["layer_rope_theta"]) | |
| assert all(t == "sparse" for t in hparams["mlp_layer_types"]) | |
| self.gguf_writer.add_sliding_window(hparams["sliding_window"]) | |
| self.gguf_writer.add_rope_freq_base_swa(theta) | |
| self.gguf_writer.add_sliding_window_pattern(sliding) | |
| # Experts live in a latent space narrower than the residual stream. | |
| self.gguf_writer.add_moe_latent_size(hparams["moe_hidden_size"]) | |
| self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) | |
| self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"]) | |
| self.gguf_writer.add_expert_weights_norm(True) | |
| # the config carries no score function, the reference hardcodes sigmoid | |
| self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) | |
| def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: | |
| if name.endswith(OFFSET_NORM_SUFFIXES): | |
| data_torch = data_torch + 1.0 | |
| if name.endswith(".self_attn.q_proj.weight"): | |
| assert bid is not None | |
| n_head = self.hparams["num_attention_heads"] | |
| head_dim = self.hparams["head_dim"] | |
| yield self.map_tensor_name(name), _unpermute_for_rope(data_torch, n_head) | |
| # Synthesize QK-norm to carry the MetaP scale so the graph keeps the plain 1 / sqrt(head_dim) | |
| ones = torch.ones(head_dim, dtype=torch.float32) | |
| q_norm = ones * self.hparams["qk_scale_factor"] | |
| yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q_NORM, bid), q_norm | |
| yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_NORM, bid), ones | |
| return | |
| if name.endswith(".self_attn.k_proj.weight"): | |
| n_head_kv = self.hparams["num_key_value_heads"] | |
| yield self.map_tensor_name(name), _unpermute_for_rope(data_torch, n_head_kv) | |
| return | |
| # gate_up_proj fuses both projections over the output dim, gate first. | |
| if name.endswith(".mlp.experts.gate_up_proj"): | |
| assert bid is not None | |
| n_ff = data_torch.shape[1] // 2 | |
| yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), data_torch[:, :n_ff, :] | |
| yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), data_torch[:, n_ff:, :] | |
| return | |
| if name.endswith(".mlp.experts.down_proj"): | |
| assert bid is not None | |
| yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), data_torch | |
| return | |
| # The router scales its logits by the MetaP output multiplier before the sigmoid | |
| # Fold it into the weight to keep a plain matmul | |
| if name.endswith(".mlp.gate.weight"): | |
| data_torch = data_torch.float() * self.hparams["output_multiplier"] | |
| yield from super().modify_tensors(data_torch, name, bid) | |
| return | |
| # Norm preserving residual gate: alpha scales the stream, beta the branch, with alpha^2 + beta^2 = 1 | |
| # Both follow from the weight, so resolve them here | |
| if name.endswith("_residual_gate.gate"): | |
| assert bid is not None | |
| attn = "post_attention" in name | |
| w = data_torch.float() / self.hparams["residual_gate_temperature"] | |
| beta = torch.sigmoid(w) | |
| alpha_sq = torch.sigmoid(-w) * (1.0 + beta) | |
| alpha = torch.sqrt(torch.clamp(alpha_sq, min=RESIDUAL_GATE_EPS)) | |
| t_alpha = gguf.MODEL_TENSOR.ATTN_RESID_ALPHA if attn else gguf.MODEL_TENSOR.FFN_RESID_ALPHA | |
| t_beta = gguf.MODEL_TENSOR.ATTN_RESID_BETA if attn else gguf.MODEL_TENSOR.FFN_RESID_BETA | |
| yield self.format_tensor_name(t_alpha, bid), alpha | |
| yield self.format_tensor_name(t_beta, bid), beta | |
| return | |
| yield from super().modify_tensors(data_torch, name, bid) | |
| class MuseSparkVisionModel(MuseGlimmerVisionModel): | |
| has_vision_encoder = True # same vision encoder as Muse Glimmer | |
| has_audio_encoder = True | |
| def get_audio_config(self) -> dict[str, Any]: | |
| c = self.global_config["audio_config"] | |
| # the base class wants one width and one layer count: use the last stage width and the total number of conv blocks | |
| return { | |
| **c, | |
| "hidden_size": c["hidden_sizes"][-1], | |
| "intermediate_size": int(c["hidden_sizes"][-1] * c["mlp_ratio"]), | |
| "num_attention_heads": 1, | |
| "num_hidden_layers": sum(c["depths"]), | |
| } | |
| def set_gguf_parameters(self): | |
| super().set_gguf_parameters() | |
| assert self.hparams_audio | |
| a = self.hparams_audio | |
| self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.MUSE_SPARK_AUDIO) | |
| # the input is the raw waveform, one channel takes the place of the mel bins | |
| self.gguf_writer.add_audio_num_mel_bins(1) | |
| self.gguf_writer.add_audio_attention_layernorm_eps(float(a["rms_norm_eps"])) | |
| self.gguf_writer.add_audio_conv_kernel_size(int(a["conv_kernel_size"])) | |
| self.gguf_writer.add_audio_hidden_sizes([int(x) for x in a["hidden_sizes"]]) | |
| self.gguf_writer.add_audio_patch_sizes([int(x) for x in a["patch_sizes"]]) | |
| self.gguf_writer.add_audio_depths([int(x) for x in a["depths"]]) | |
| def filter_tensors(cls, item): | |
| name, gen = item | |
| if name.startswith("model.audio_tower."): | |
| return MmprojModel.filter_tensors((name, gen)) | |
| return super().filter_tensors((name, gen)) | |
| _AUDIO_HEAD = { | |
| "model.audio_tower.norm": gguf.MODEL_TENSOR.A_PRE_NORM, | |
| "model.audio_tower.proj": gguf.MODEL_TENSOR.A_MMPROJ_FC, | |
| "model.audio_tower.post_proj_norm": gguf.MODEL_TENSOR.A_POST_NORM, | |
| } | |
| def _audio_tensor(self, data_torch, name, suffix): | |
| assert self.hparams_audio | |
| stem = name[: -len(suffix) - 1] | |
| if stem in self._AUDIO_HEAD: | |
| return self.format_tensor_name(self._AUDIO_HEAD[stem], suffix="." + suffix), data_torch | |
| parts = stem.split(".") | |
| stage = int(parts[3]) | |
| if parts[4] == "patch_embedder": | |
| key = (gguf.MODEL_TENSOR.A_ENC_PATCH_PROJ if parts[5] == "proj" | |
| else gguf.MODEL_TENSOR.A_ENC_PATCH_NORM) | |
| return self.format_tensor_name(key, bid=stage, suffix="." + suffix), data_torch | |
| # conv blocks are numbered per stage upstream and flat here | |
| depths = [int(d) for d in self.hparams_audio["depths"]] | |
| bid = sum(depths[:stage]) + int(parts[5]) | |
| tail = ".".join(parts[6:]) | |
| key = { | |
| "conv": gguf.MODEL_TENSOR.A_ENC_CONV1D, | |
| "norm": gguf.MODEL_TENSOR.A_ENC_CONV1D_NORM, | |
| "mlp.fc1": gguf.MODEL_TENSOR.A_ENC_FFN_UP, | |
| "mlp.fc2": gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, | |
| }[tail] | |
| return self.format_tensor_name(key, bid=bid, suffix="." + suffix), data_torch | |
| def modify_tensors(self, data_torch, name, bid): | |
| if name.startswith("model.audio_tower."): | |
| yield self._audio_tensor(data_torch, name, name.rpartition(".")[2]) | |
| return | |
| yield from super().modify_tensors(data_torch, name, bid) | |
Xet Storage Details
- Size:
- 9.87 kB
- Xet hash:
- e37d856afc0138f1e87fa9c51fdf97f462fa454f9ca7ec0767851eac7756d38d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.