Buckets:
| from __future__ import annotations | |
| from typing import Iterable, cast | |
| import torch | |
| from torch import Tensor | |
| import gguf | |
| import numpy as np | |
| from .base import ModelBase | |
| from .qwen import _LinearAttentionVReorderBase, _Qwen35MRopeMixin | |
| from .qwen3vl import Qwen3VLVisionModel | |
| class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): | |
| """Qwen3.8-Flash-Next. | |
| Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things: | |
| hyper-connections in place of every layer norm, QSA sparse attention on the full | |
| attention layers, and PLE n-gram hash embeddings on a single layer. | |
| """ | |
| model_arch = gguf.MODEL_ARCH.QWEN4EXP | |
| # the MTP head: one full-attention QSA block after the trunk, fed by the trunk's hc-wide residual | |
| supports_mtp_export = True | |
| # MTP tensors the shared Qwen remapper does not know | |
| _MTP_EXTRA = { | |
| "fc_embedding": "nextn_fc_embedding", | |
| "fc_hidden": "nextn_fc_hidden", | |
| "hyper_connection_mixer": "nextn_hc_head", | |
| } | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| # only the shard names, so the table itself is never held | |
| self._ple_shards: dict[int, str] = {} | |
| self._ple_row_dim: int | None = None | |
| self._mtp_fc: dict[str, Tensor] = {} | |
| # halves of the merged weights, keyed by (bid, merged tensor), see _merge_shared_input | |
| self._merge_parts: dict[tuple[int, gguf.MODEL_TENSOR], dict[int, Tensor]] = {} | |
| def filter_tensors(cls, item): | |
| name, gen = item | |
| part = name.split(".")[1] if name.startswith("mtp.") else None | |
| if part in cls._MTP_EXTRA: | |
| if cls.no_mtp: | |
| return None | |
| assert cls._original_block_count is not None | |
| rest = name.split(".", 2)[2] | |
| return f"model.layers.{cls._original_block_count}.{cls._MTP_EXTRA[part]}.{rest}", gen | |
| return super().filter_tensors(item) | |
| def _read_hash_constants(self, suffix: str) -> list[int]: | |
| """Read an int64 PLE constant straight from the checkpoint. | |
| prepare_tensors() casts every non-float dtype to float32 before | |
| modify_tensors() sees it (base.py), which would silently round these | |
| 45-bit multipliers. Reading the lazy tensor here bypasses that. | |
| """ | |
| for name, gen in self.model_tensors.items(): | |
| if name.endswith(suffix): | |
| t = gen() | |
| if t.dtype != torch.int64: | |
| t = t.to(torch.int64) | |
| return [int(x) for x in t.tolist()] | |
| raise ValueError(f"PLE constant {suffix!r} missing from the checkpoint") | |
| def set_gguf_parameters(self): | |
| super().set_gguf_parameters() | |
| hp = self.hparams | |
| self.gguf_writer.add_hyper_connection_count(hp["hc_count"]) | |
| self.gguf_writer.add_hyper_connection_low_rank(hp["hc_lowrank"]) | |
| n_layer = hp["num_hidden_layers"] | |
| self.gguf_writer.add_indexer_head_count(hp["indexer_n_heads"]) | |
| self.gguf_writer.add_indexer_key_length(hp["indexer_head_dim"]) | |
| self.gguf_writer.add_indexer_top_k(hp["indexer_budget"]) | |
| ratio = hp["indexer_compress_ratio"] | |
| layer_types = hp["layer_types"] | |
| # the MTP block is a full-attention QSA layer too | |
| self.gguf_writer.add_attention_compress_ratios( | |
| [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] | |
| + [ratio] * (self.block_count - n_layer) | |
| ) | |
| # ple_layer_ids is 1-based in the HF config; empty means no n-gram table, | |
| # so emit no PLE keys rather than optional ones | |
| # the MTP head never reads PLE, so an MTP-only file carries none of it | |
| ple_layers = [i - 1 for i in hp["ple_layer_ids"]] | |
| if not ple_layers or self.mtp_only: | |
| return | |
| self.gguf_writer.add_ple_layers(ple_layers) | |
| self.gguf_writer.add_ple_ngram_size(hp["ngram_size"]) | |
| self.gguf_writer.add_ple_heads_per_ngram(hp["heads_per_ngram"]) | |
| self.gguf_writer.add_ple_conv_kernel(hp["ple_conv_kernel_size"]) | |
| self.gguf_writer.add_ple_eos_token_id(self._eos_token_id()) | |
| # an image is decoded as an embeddings-only batch, so the graph has no placeholder | |
| # ids to hash; carry the id and let it stand in for those positions | |
| _img = self._image_token_id() | |
| if _img is not None: | |
| self.gguf_writer.add_ple_image_token_id(int(_img)) | |
| if self._ple_row_dim is not None: | |
| self.gguf_writer.add_embedding_length_per_layer_input(self._ple_row_dim) | |
| self.gguf_writer.add_ple_layer_multipliers( | |
| self._read_hash_constants("ple_embedding.layer_multipliers")) | |
| self.gguf_writer.add_ple_head_offsets( | |
| self._read_hash_constants("ple_embedding.ngram_heads_offsets")) | |
| self.gguf_writer.add_ple_head_vocab_sizes( | |
| self._read_hash_constants("ple_embedding.ngram_heads_vocab_sizes")) | |
| def _image_token_id(self) -> int | None: | |
| img = self.hparams.get("image_token_id") | |
| return None if img is None else int(img) | |
| def _eos_token_id(self) -> int: | |
| eos = self.hparams.get("eos_token_id") | |
| if isinstance(eos, list): | |
| # the PLE hash resets n-grams on the primary EOS | |
| return int(eos[-1]) | |
| if eos is None: | |
| raise ValueError("eos_token_id is required: the PLE hash resets its n-grams on it") | |
| return int(eos) | |
| def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: | |
| # int64 hash constants must stay exact; 1-D tensors force F32, so use KV | |
| if name.endswith("ple_embedding.layer_multipliers"): | |
| self._ple_multipliers = [int(x) for x in data_torch.tolist()] | |
| return | |
| if name.endswith("ple_embedding.ngram_heads_offsets"): | |
| self._ple_head_offsets = [int(x) for x in data_torch.tolist()] | |
| return | |
| if name.endswith("ple_embedding.ngram_heads_vocab_sizes"): | |
| self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()] | |
| return | |
| if ".ngram_embedding.shard_" in name: | |
| yield from self._place_ple_shard(data_torch, name) | |
| return | |
| # eh_proj([e ; h_s]) = fc_embedding(e) + fc_hidden(h_s) for every hc stream s | |
| if name.endswith((".nextn_fc_embedding.weight", ".nextn_fc_hidden.weight")): | |
| self._mtp_fc[name.rsplit(".", 2)[1]] = data_torch | |
| if len(self._mtp_fc) < 2: | |
| return | |
| eh = torch.cat([self._mtp_fc.pop("nextn_fc_embedding"), self._mtp_fc.pop("nextn_fc_hidden")], dim=1) | |
| yield (self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh) | |
| return | |
| # one projection feeds indexer q and k; split it, as minimax-m3 does | |
| if ".indexer.index_qk_proj.weight" in name: | |
| n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"] | |
| yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), data_torch[:n_q]) | |
| yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), data_torch[n_q:]) | |
| return | |
| # Gemma zero-centred gammas the inherited norm.weight rule misses | |
| if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight", | |
| ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")): | |
| yield (self.map_tensor_name(name), data_torch + 1) | |
| return | |
| if name.endswith(".ple.conv1d.weight"): | |
| yield (self.map_tensor_name(name), data_torch.squeeze()) | |
| return | |
| for new_name, data in super().modify_tensors(data_torch, name, bid): | |
| yield from self._merge_shared_input(new_name, data, bid) | |
| # matmuls that read the same input, written as one weight [a | b] so the runtime does one matmul: | |
| # (merged tensor, (a, b), linear-attention layers only) | |
| _SHARED_INPUT_MERGES = ( | |
| (gguf.MODEL_TENSOR.HC_ATTN_DOWN_INJECT, (gguf.MODEL_TENSOR.HC_ATTN_DOWN, gguf.MODEL_TENSOR.HC_ATTN_INJECT), False), | |
| (gguf.MODEL_TENSOR.HC_FFN_DOWN_INJECT, (gguf.MODEL_TENSOR.HC_FFN_DOWN, gguf.MODEL_TENSOR.HC_FFN_INJECT), False), | |
| (gguf.MODEL_TENSOR.ATTN_QKV_GATE, (gguf.MODEL_TENSOR.ATTN_QKV, gguf.MODEL_TENSOR.ATTN_GATE), True), | |
| (gguf.MODEL_TENSOR.SSM_BETA_ALPHA_CAT, (gguf.MODEL_TENSOR.SSM_BETA, gguf.MODEL_TENSOR.SSM_ALPHA), True), | |
| ) | |
| # caches one half of a merged weight until the other half arrives, then yields [a | b] | |
| def _merge_shared_input(self, new_name: str, data: Tensor, bid: int | None) -> Iterable[tuple[str, Tensor]]: | |
| # a pre-quantized checkpoint carries per-tensor scales, which one merged weight cannot hold | |
| if bid is None or self.hparams.get("quantization_config"): | |
| yield (new_name, data) | |
| return | |
| layer_types = self.hparams.get("layer_types") or [] | |
| is_linear = bid < len(layer_types) and layer_types[bid] == "linear_attention" | |
| match = next(((merged, parts.index(part)) for merged, parts, linear_only in self._SHARED_INPUT_MERGES | |
| for part in parts | |
| if (is_linear or not linear_only) and self.match_model_tensor_name(new_name, part, bid)), None) | |
| if match is None: | |
| yield (new_name, data) | |
| return | |
| merged, idx = match | |
| halves = self._merge_parts.setdefault((bid, merged), {}) | |
| halves[idx] = data | |
| if len(halves) == 2: | |
| del self._merge_parts[(bid, merged)] | |
| yield (self.format_tensor_name(merged, bid), torch.cat([halves[0], halves[1]], dim=0)) | |
| # the shards concatenate into a tensor of well over 100 GB | |
| # use LazyChunkedTensor here, a single shard resident at a time | |
| def _place_ple_shard(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]: | |
| idx = int(name.rpartition(".shard_")[2].partition(".")[0]) | |
| n_parts = self.hparams["split_ngram_parts"] | |
| self._ple_shards[idx] = name | |
| self._ple_row_dim = int(data_torch.shape[-1]) | |
| if len(self._ple_shards) < n_parts: | |
| return [] | |
| # the checkpoint may yield the shards in any order, the row order is by index | |
| shards = [self._ple_shards[i] for i in sorted(self._ple_shards)] | |
| rows = 0 | |
| for shard in shards: | |
| shape = self.model_tensors[shard]().shape | |
| if int(shape[-1]) != self._ple_row_dim: | |
| raise ValueError( | |
| f"PLE shard {shard} has row dim {int(shape[-1])}, expected {self._ple_row_dim}") | |
| rows += int(shape[0]) | |
| table = gguf.LazyChunkedTensor( | |
| [self._load_ple_shard(shard) for shard in shards], | |
| shape=(rows, self._ple_row_dim), | |
| dtype=np.float32, | |
| ) | |
| gguf_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD] | |
| return [(gguf_name + ".weight", cast(Tensor, table))] | |
| def _load_ple_shard(self, name: str): | |
| def load() -> np.ndarray: | |
| from .base import LazyTorchTensor | |
| # a fresh lazy tensor every call, or to_eager() memoizes every shard | |
| eager = LazyTorchTensor.to_eager(self.model_tensors[name]()) | |
| return eager.to(torch.float32).contiguous().numpy() | |
| return load | |
| def prepare_tensors(self): | |
| super().prepare_tensors() | |
| if self._mtp_fc: | |
| raise ValueError(f"MTP projection missing its other half: {sorted(self._mtp_fc)}") | |
| if self._merge_parts: | |
| raise ValueError(f"merged weights missing a half: {sorted((bid, t.name) for bid, t in self._merge_parts)}") | |
| n_parts = self.hparams.get("split_ngram_parts", 0) | |
| if self._ple_shards and len(self._ple_shards) != n_parts: | |
| raise ValueError( | |
| f"got {len(self._ple_shards)} PLE embedding shards, expected {n_parts}" | |
| ) | |
| class Qwen4ExpVisionModel(Qwen3VLVisionModel): | |
| """The vision tower is an unmodified Qwen3-VL ViT.""" | |
Xet Storage Details
- Size:
- 12.4 kB
- Xet hash:
- a366b8ca7f72b7f16372344c226b961f7e9e219899e981b828139b9f0a9c6811
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.