From 02d76b655011e20aea7c673d6117885c0884fdc3 Mon Sep 17 00:00:00 2001 From: Mukkesh Ganesh Date: Tue, 20 Oct 2026 16:30:57 -0700 Subject: [PATCH 1/3] feat(gguf): serve DeepSeek-V4-Flash GGUF checkpoints on deepseek4 --- python/freetoken/kernel/aot_models.py | 2 + .../freetoken/models/deepseek_v4/__init__.py | 7 +- python/freetoken/models/deepseek_v4/config.py | 51 +-- python/freetoken/models/deepseek_v4/gguf.py | 296 ++++++++++++++++++ python/freetoken/models/deepseek_v4/model.py | 2 +- python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/gguf/tokenizer.py | 15 + python/freetoken/models/register.py | 20 +- 8 files changed, 367 insertions(+), 27 deletions(-) create mode 100644 python/freetoken/models/deepseek_v4/gguf.py diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index a68c868ea7..e45b98d5c9 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -326,6 +326,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int AotModel( name="deepseek-ai/DeepSeek-V4-Flash", architecture="DeepseekV4ForCausalLM", + # llama.cpp deepseek4 GGUF (unsloth 0731 GGUF; text only) + arch_aliases=("DeepseekV4GGUFForCausalLM",), hidden_size=4096, kv_groups=(), # MLA latent pool writes via torch scatter, not store_cache top_k=6, diff --git a/python/freetoken/models/deepseek_v4/__init__.py b/python/freetoken/models/deepseek_v4/__init__.py index 5bbb9b8e72..31931ede49 100644 --- a/python/freetoken/models/deepseek_v4/__init__.py +++ b/python/freetoken/models/deepseek_v4/__init__.py @@ -14,7 +14,8 @@ """ from .args import DeepseekV4Args, load_args -from .config import DSV4VisionConfig, parse_config, parse_vision_config +from .config import DSV4VisionConfig, model_config_from_args, parse_config, parse_vision_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import DeepseekV4ForCausalLM from .vision import DSV4Aligner, DSV4Vision, DSV4VisionTower from .weight import iter_expert_pieces, iter_vision_weights, iter_weights @@ -26,10 +27,14 @@ "DSV4VisionConfig", "DSV4VisionTower", "load_args", + "model_config_from_args", "parse_config", + "parse_gguf_config", "parse_vision_config", "DeepseekV4ForCausalLM", + "iter_gguf_weights", "iter_weights", "iter_expert_pieces", "iter_vision_weights", + "to_gguf", ] diff --git a/python/freetoken/models/deepseek_v4/config.py b/python/freetoken/models/deepseek_v4/config.py index 9671c68290..65196a4d66 100644 --- a/python/freetoken/models/deepseek_v4/config.py +++ b/python/freetoken/models/deepseek_v4/config.py @@ -18,7 +18,7 @@ from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig -from .args import load_args +from .args import DeepseekV4Args, load_args # ``<｜deepseek_image｜>``: the only vocab-level image token. The reference expands each # occurrence into a block of ``vocab_size + type`` pseudo ids; here the block keeps this @@ -46,16 +46,12 @@ class DSV4VisionConfig: text_dim: int -def parse_config(hf_config: Any) -> ModelConfig: - model_path = getattr(hf_config, "_name_or_path", None) or getattr( - hf_config, "name_or_path", None - ) - if not model_path: - raise ValueError( - "DeepSeek-V4 parse_config needs the checkpoint path (hf_config._name_or_path)" - ) - args = load_args(model_path, max_batch_size=1) - +def model_config_from_args( + args: DeepseekV4Args, *, max_position: int, vision_config: DSV4VisionConfig | None = None +) -> ModelConfig: + """The family's :class:`ModelConfig` from built :class:`DeepseekV4Args`: the one + construction the HF checkpoint parser and the GGUF parser share, so the two flows + cannot drift. ``expert_quant`` is the caller's to override (``ds_fp4`` vs ``gguf``).""" rope_scaling = { "rope_type": "yarn", "factor": args.rope_factor, @@ -64,16 +60,6 @@ def parse_config(hf_config: Any) -> ModelConfig: "original_max_position_embeddings": args.original_seq_len, } - # Serving ceiling: the HF top-level config's max_position_embeddings (1M on V4-Flash). - # args.original_seq_len is yarn's PRE-scaling length (64k); the compressed tiers serve - # original_seq_len * rope_factor positions, so fall back to that product when a checkpoint - # ships only the inference/ dialect. - max_position = int(getattr(hf_config, "max_position_embeddings", 0) or 0) - if max_position <= 0: - max_position = int(args.original_seq_len * args.rope_factor) - - vision_config = parse_vision_config(hf_config) - return ModelConfig( num_layers=args.n_layers, num_qo_heads=args.n_heads, @@ -124,6 +110,27 @@ def parse_config(hf_config: Any) -> ModelConfig: ) +def parse_config(hf_config: Any) -> ModelConfig: + model_path = getattr(hf_config, "_name_or_path", None) or getattr( + hf_config, "name_or_path", None + ) + if not model_path: + raise ValueError( + "DeepSeek-V4 parse_config needs the checkpoint path (hf_config._name_or_path)" + ) + args = load_args(model_path, max_batch_size=1) + + # Serving ceiling: the HF top-level config's max_position_embeddings (1M on V4-Flash). + # args.original_seq_len is yarn's PRE-scaling length (64k); the compressed tiers serve + # original_seq_len * rope_factor positions, so fall back to that product when a checkpoint + # ships only the inference/ dialect. + max_position = int(getattr(hf_config, "max_position_embeddings", 0) or 0) + if max_position <= 0: + max_position = int(args.original_seq_len * args.rope_factor) + + return model_config_from_args(args, max_position=max_position, vision_config=parse_vision_config(hf_config)) + + def parse_vision_config(hf_config: Any) -> DSV4VisionConfig | None: """The tower's dims from the checkpoint config, or None when it carries no vision stack. @@ -150,4 +157,4 @@ def parse_vision_config(hf_config: Any) -> DSV4VisionConfig | None: ) -__all__ = ["IMAGE_TOKEN_ID", "DSV4VisionConfig", "parse_config", "parse_vision_config"] +__all__ = ["IMAGE_TOKEN_ID", "DSV4VisionConfig", "model_config_from_args", "parse_config", "parse_vision_config"] diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py new file mode 100644 index 0000000000..50457fcd86 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -0,0 +1,296 @@ +"""DeepSeek-V4-Flash from a llama.cpp ``deepseek4`` GGUF (unsloth ``DeepSeek-V4-Flash-0731-GGUF``, +text only): config from the GGUF metadata, tensor names, and the dense state dict. + +The GGUF geometry matches the HF checkpoint (43 MLA blocks with sliding-window sparse +attention, per-layer KV compressors and a Lightning Indexer on the ratio-4 blocks, mHC x4 +residual streams, a 256-expert sqrtsoftplus MoE with hash routing on the first blocks and +a shared expert), so ``parse_gguf_config`` builds the family's own :class:`DeepseekV4Args` +from the GGUF KV metadata and reuses the shared ModelConfig construction +(:func:`deepseek_v4.config.model_config_from_args`). + +The converter drops the MTP blocks and writes ``attention.compress_ratios`` as a single +placeholder 0, so the per-layer compressor layout is recovered from the tensor table: an +``indexer.*`` tensor means ratio 4, a compressor without an indexer means 128, unless its +``wkv`` carries the CSA overlap's double rows (also ratio 4). Routed experts serve from +the GGUF expert banks (``freetoken.models.gguf.experts``) like every GGUF family. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from typing import TYPE_CHECKING, Iterator + +import numpy as np +import torch + +from freetoken.layers.quantization.ggml_types import is_quantized +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index, gguf_tensor_names +from freetoken.models.gguf.tensors import check_scales, fused, matrix + +from .args import DeepseekV4Args +from .config import model_config_from_args + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +# engine module leaf (after "model.layers.N.") -> GGUF tensor names (without ".weight") +_LAYER_LEAVES: dict[str, tuple[str, ...]] = { + "attn.wq_a": ("attn_q_a",), + "attn.wq_b": ("attn_q_b",), + "attn.wkv": ("attn_kv",), + "attn.wo_b": ("attn_output_b",), + "attn.compressor.wkv": ("attn_compressor_kv",), + "attn.compressor.wgate": ("attn_compressor_gate",), + "attn.indexer.wq_b": ("indexer.attn_q_b",), + "attn.indexer.weights_proj": ("indexer.proj",), + "attn.indexer.compressor.wkv": ("indexer_compressor_kv",), + "attn.indexer.compressor.wgate": ("indexer_compressor_gate",), + "ffn.shared_experts.w1": ("ffn_gate_shexp",), + "ffn.shared_experts.w3": ("ffn_up_shexp",), + "ffn.shared_experts.w2": ("ffn_down_shexp",), + # routed experts (GGUF expert banks), gate / up / down order + "ffn.experts.0.w1": ("ffn_gate_exps",), + "ffn.experts.0.w3": ("ffn_up_exps",), + "ffn.experts.0.w2": ("ffn_down_exps",), +} +_TOP_LEAVES: dict[str, tuple[str, ...]] = { + "model.embed": ("token_embd",), + "model.head": ("output",), +} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Engine module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return _TOP_LEAVES[name] + m = _LAYER_RE.match(name) + if m is None: + return () + return tuple(f"blk.{m.group(1)}.{leaf}" for leaf in _LAYER_LEAVES.get(m.group(2), ())) + + +def _compress_ratios(index: dict[str, GgufTensorInfo], head_dim: int, num_layers: int) -> tuple[int, ...]: + """Per-layer compressor ratios from the tensor table: the KV placeholder carries a single 0. + + An ``indexer.proj`` tensor exists exactly on the ratio-4 blocks; a compressor without an + indexer is a 128 block, unless its ``wkv`` has the CSA overlap's double rows (ratio 4).""" + ratios = [] + for layer in range(num_layers): + blk = f"blk.{layer}" + if f"{blk}.indexer.proj.weight" in index: + ratios.append(4) + continue + info = index.get(f"{blk}.attn_compressor_kv.weight") + ratios.append(0 if info is None else 4 if info.ne[-1] == 2 * head_dim else 128) + return tuple(ratios) + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + md = shim.metadata + prefix = shim.model_type + + check_scales(shim.model_path) + + def kv(key: str, default=None): + return md.get(f"{prefix}.{key}", default) + + def need(key: str): + val = kv(key) + if val is None: + raise KeyError(f"missing GGUF metadata key {prefix}.{key}") + return val + + names = gguf_tensor_names(shim.model_path) + if any(n.startswith(("vision.", "aligner.")) for n in names): + raise NotImplementedError( + f"{shim.model_path}: deepseek4 GGUF with a vision tower is not served (the text-only 0731 release has none)" + ) + + # the clamp is per-layer in the GGUF (and the shared expert carries its own array); the model + # takes one value, so anything non-uniform is refused instead of silently taking the first + clamp = need("swiglu_clamp_exp") + clamp_vals = list(clamp) if isinstance(clamp, list) else [clamp] + clamp_shexp = kv("swiglu_clamp_shexp") + shexp_vals = list(clamp_shexp) if isinstance(clamp_shexp, list) else [clamp_shexp] if clamp_shexp is not None else [] + if len(set(map(float, clamp_vals + shexp_vals))) != 1: + raise NotImplementedError(f"{shim.model_path}: non-uniform swiglu_clamp_exp / swiglu_clamp_shexp is not served") + + from freetoken.third_party.gguf.constants import ExpertGatingFuncType + + if int(need("expert_gating_func")) != int(ExpertGatingFuncType.SQRTSOFTPLUS): + raise NotImplementedError(f"{shim.model_path}: expert_gating_func other than sqrtsoftplus is not served") + + if kv("rope.scaling.type", "yarn") != "yarn": + raise NotImplementedError(f"{shim.model_path}: rope scaling other than yarn is not served") + + head_dim = int(need("attention.key_length")) + if int(need("attention.value_length")) != head_dim: + raise NotImplementedError(f"{shim.model_path}: attention key and value lengths differ") + if int(kv("attention.head_count_kv", 1) or 1) != 1: + raise NotImplementedError(f"{shim.model_path}: more than one latent KV head is not served") + + num_layers = int(need("block_count")) + index = gguf_tensor_index(shim.model_path) + args = DeepseekV4Args( + max_batch_size=1, + vocab_size=int(shim.vocab_size), + dim=int(need("embedding_length")), + moe_inter_dim=int(need("expert_feed_forward_length")), + n_layers=num_layers, + n_hash_layers=int(kv("hash_layer_count", 0) or 0), + n_heads=int(need("attention.head_count")), + n_routed_experts=int(need("expert_count")), + n_shared_experts=int(kv("expert_shared_count", 1) or 1), + n_activated_experts=int(need("expert_used_count")), + score_func="sqrtsoftplus", + route_scale=float(need("expert_weights_scale")), + swiglu_limit=float(clamp_vals[0]), + q_lora_rank=int(need("attention.q_lora_rank")), + head_dim=head_dim, + rope_head_dim=int(need("rope.dimension_count")), + norm_eps=float(need("attention.layer_norm_rms_epsilon")), + o_groups=int(need("attention.output_group_count")), + o_lora_rank=int(need("attention.output_lora_rank")), + window_size=int(need("attention.sliding_window")), + compress_ratios=_compress_ratios(index, head_dim, num_layers), + compress_rope_theta=float(kv("attention.compress_rope_freq_base", 160000.0) or 160000.0), + original_seq_len=int(need("rope.scaling.original_context_length")), + rope_theta=float(need("rope.freq_base")), + rope_factor=float(need("rope.scaling.factor")), + beta_fast=float(kv("rope.scaling.yarn_beta_fast", 32) or 32), + beta_slow=float(kv("rope.scaling.yarn_beta_slow", 1) or 1), + index_n_heads=int(need("attention.indexer.head_count")), + index_head_dim=int(need("attention.indexer.key_length")), + index_topk=int(need("attention.indexer.top_k")), + hc_mult=int(need("hyper_connection.count")), + hc_sinkhorn_iters=int(need("hyper_connection.sinkhorn_iterations")), + hc_eps=float(need("hyper_connection.epsilon")), + n_mtp_layers=0, # the converter drops the MTP blocks + vision_n_layers=0, + ) + config = model_config_from_args(args, max_position=int(need("context_length"))) + return replace(config, expert_quant="gguf") + + +# ---------------------------------------------------------------------------------------------- +# Weights +# ---------------------------------------------------------------------------------------------- + +def _float(info: GgufTensorInfo) -> torch.Tensor: + """A float tensor in torch order, dequantizing packed ggml rows: dynamic quants quantize + tensors the module consumes as floats (the mHC mixers, the compressor gates, the router).""" + if not is_quantized(info.ggml_type): + from freetoken.models.gguf.tensors import float_tensor + + return float_tensor(info) + from freetoken.third_party.gguf import quants + from freetoken.third_party.gguf.constants import GGMLQuantizationType + + flat = quants.dequantize(np.array(info.data()), GGMLQuantizationType(info.ggml_type)) + return torch.from_numpy(flat).view(*reversed(info.ne)) + + +def _tid2eid(info: GgufTensorInfo) -> torch.Tensor: + """The hash-route table as the module's int64 [vocab, top_k].""" + raw = np.frombuffer(info.data(), dtype=np.int32).copy() + return torch.from_numpy(raw).view(*reversed(info.ne)).to(torch.int64) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, + include_vision: bool = True, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF, in the names the HF loader yields; the routed experts + load as offload banks (``freetoken.models.gguf.experts``).""" + if include_moe_experts: + raise ValueError( + "DeepSeek-V4 routed experts are served from the offload cache; " + "run with --moe-strategy offload (include_moe_experts must be False)." + ) + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + config = parse_gguf_config(cached_load_hf_config(model_path)) + args = config.dsv4_args + index = gguf_tensor_index(model_path) + hidden = config.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def bf16(name: str) -> torch.Tensor: + return _float(t(name)).to(torch.bfloat16) + + yield from fused("model.embed", [t("token_embd")], [matrix(t("token_embd"))], hidden) + yield from fused("model.head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", bf16("output_norm") + for part in ("fn", "base", "scale"): + yield f"model.hc_head_{part}", _float(t(f"output_hc_{part}")) + + for layer in range(config.num_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + m = f"{pre}.attn" + yield f"{pre}.attn_norm.weight", bf16(f"{blk}.attn_norm") + yield f"{pre}.ffn_norm.weight", bf16(f"{blk}.ffn_norm") + # mHC mixing tensors, fp32 on every layer + for sub in ("attn", "ffn"): + for part in ("fn", "base", "scale"): + yield f"{pre}.hc_{sub}_{part}", _float(t(f"{blk}.hc_{sub}_{part}")) + + for proj, src, k in ( + ("wq_a", "attn_q_a", hidden), + ("wq_b", "attn_q_b", args.q_lora_rank), + ("wkv", "attn_kv", hidden), + ("wo_b", "attn_output_b", args.o_groups * args.o_lora_rank), + ): + info = t(f"{blk}.{src}") + yield from fused(f"{m}.{proj}", [info], [matrix(info)], k) + yield f"{m}.q_norm.weight", bf16(f"{blk}.attn_q_a_norm") + yield f"{m}.kv_norm.weight", bf16(f"{blk}.attn_kv_a_norm") + # wo_a feeds a bf16 einsum; the HF path dequantizes its fp8 the same way + yield f"{m}.wo_a", _float(t(f"{blk}.attn_output_a")).to(torch.bfloat16) + yield f"{m}.attn_sink", _float(t(f"{blk}.attn_sinks")) + + ratio = args.compress_ratios[layer] + if ratio: + compressors = ((f"{m}.compressor", f"{blk}.attn_compressor"),) + if ratio == 4: + compressors += ((f"{m}.indexer.compressor", f"{blk}.indexer_compressor"),) + for c_dst, c_src in compressors: + yield f"{c_dst}.ape", _float(t(f"{c_src}_ape")) + # the compressors read wkv / wgate through a dense bf16 einsum, not the quant + # method (the registry exempts them, like the HF spec), so a tensor the file + # quantized dequantizes here + for proj, suf in (("wkv", "kv"), ("wgate", "gate")): + yield f"{c_dst}.{proj}.weight", bf16(f"{c_src}_{suf}") + yield f"{c_dst}.norm.weight", bf16(f"{c_src}_norm") + if ratio == 4: + for proj, src, k in ( + ("indexer.wq_b", "indexer.attn_q_b", args.q_lora_rank), + ("indexer.weights_proj", "indexer.proj", hidden), + ): + info = t(f"{blk}.{src}") + yield from fused(f"{m}.{proj}", [info], [matrix(info)], k) + + g = f"{pre}.ffn.gate" + yield f"{g}.weight", bf16(f"{blk}.ffn_gate_inp") + if layer < args.n_hash_layers: + # a vision checkpoint also gives its hash layers a text bias; the text model has no slot for it + yield f"{g}.tid2eid", _tid2eid(t(f"{blk}.ffn_gate_tid2eid")) + else: + yield f"{g}.bias", _float(t(f"{blk}.exp_probs_b.bias")) + for proj, src in (("w1", "ffn_gate_shexp"), ("w3", "ffn_up_shexp"), ("w2", "ffn_down_shexp")): + info = t(f"{blk}.{src}") + k = hidden if proj != "w2" else args.moe_inter_dim + yield from fused(f"{pre}.ffn.shared_experts.{proj}", [info], [matrix(info)], k) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index 0c3441b951..77ef7f6162 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -130,7 +130,7 @@ def __init__(self, args: DeepseekV4Args, quant_config=None, *, strategy: str = " self.norm_eps = args.norm_eps self.hc_eps = args.hc_eps self.hc_mult = hc_mult = args.hc_mult - self.embed = VocabParallelEmbedding(args.vocab_size, args.dim) + self.embed = VocabParallelEmbedding(args.vocab_size, args.dim, quant_config=quant_config, prefix=f"{prefix}.embed") self.layers = OPList([Block(i, args, strategy=strategy, decode_target=decode_target, quant_config=quant_config, prefix=f"{prefix}.layers.{i}") for i in range(args.n_layers)]) self.norm = RMSNorm(args.dim, self.norm_eps) self.head = ParallelLMHead(args.vocab_size, args.dim, quant_config=quant_config, prefix=f"{prefix}.head") diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 5767d4eb1c..214fb51bcc 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -17,6 +17,7 @@ # GGUF ``general.architecture`` -> FreeToken registry key (a GGUF-specific spec that # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { + "deepseek4": "DeepseekV4GGUFForCausalLM", "gemma4": "Gemma4GGUFForCausalLM", "glm4moe": "Glm4MoeGGUFForCausalLM", # glm5-next: the llama.cpp spelling; glm5next: unsloth's first release diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 26e447aad2..5d0ff489ce 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -25,6 +25,14 @@ # Llama 3 and GLM pre-tokenizers, whose HF tokenizer and llama.cpp take a word that is one vocabulary token as is instead of merging it _IGNORE_MERGES_PRE = frozenset(("llama3", "llama-v3", "llama-bpe", "glm4", "glm5")) +# DeepSeek-V4's pre-tokenizer (llama.cpp spells it "joyai-llm"; transformers maps no split +# for it): the three splits of the model's own tokenizer.json, in order +_JOYAI_LLM_SPLITS = ( + r"\p{N}{1,3}", + "[一-龥぀-ゟ゠-ヿ]+", + r'''[!"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+| ?[\p{P}\p{S}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+''', +) + def _add_bos(meta: dict) -> bool: if "tokenizer.ggml.add_bos_token" in meta: @@ -91,6 +99,13 @@ def load_gguf_tokenizer(model_path: str): pre_tokenizers.Split(Regex(GGUF_PRE_TOKENIZER_SPLITS["llama3"]), behavior="isolated"), pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False), ]) + elif pre == "joyai-llm": + from tokenizers import Regex, pre_tokenizers + + fast.pre_tokenizer = pre_tokenizers.Sequence( + [pre_tokenizers.Split(Regex(r), behavior="isolated") for r in _JOYAI_LLM_SPLITS] + + [pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)] + ) if pre in _IGNORE_MERGES_PRE: fast.model.ignore_merges = True # special: bos/eos/unk/pad as the file names them (None when absent) and the chat template diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 4d268ffb50..65601d63db 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -208,14 +208,28 @@ class ModelSpec: # the checkpoint has no ``model.`` root checkpoint_roots=(("model.layers", "layers"), ("model.head", "head")), packed_modules_mapping=_EXPERTS_W123_PACKED, - # the head, the KV compressors, the indexer's scorer and the vision tower + aligner ship bf16; - # the fp8 config has no modules_to_not_convert (patterns match checkpoint names, like V4.1) + # the embed, the head, the KV compressors, the indexer's scorer and the vision tower + + # aligner ship bf16; the fp8 config has no modules_to_not_convert (patterns match + # checkpoint names, like V4.1) unquantized_modules=( - "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj", "vision.*", "aligner.*", + "embed", "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj", "vision.*", "aligner.*", ), mm_processor=_DSV4_PROCESSOR, encoders=_DSV4_ENCODERS, ), + # llama.cpp deepseek4 GGUF (unsloth DeepSeek-V4-Flash-0731-GGUF; text only): the same model + # classes with every tensor in its packed ggml type + "DeepseekV4GGUFForCausalLM": ModelSpec( + "freetoken.models.deepseek_v4", + "DeepseekV4ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + # the GGUF scheme reads expert types in (gate, up, down) order + packed_modules_mapping=_EXPERTS_W132_PACKED, + # the KV compressors read their weights through a dense bf16 einsum, so the reader + # dequantizes them instead of serving the packed ggml rows + unquantized_modules=("*.compressor.wkv", "*.compressor.wgate"), + ), # DeepSeek-V4.1-Flash (model_type deepseek_v41): multimodal wrapper config (text tower in # text_config). The checkpoint has no ``model.`` root; the compressors, the indexer's key projection / # scorer, the head and the vision tower + aligner ship bf16 (the fp8 config has no modules_to_not_convert). From a48447e2e90a39cc6c3d9849fcc84043deb2cc28 Mon Sep 17 00:00:00 2001 From: Mukkesh Ganesh Date: Tue, 20 Oct 2026 18:12:01 -0700 Subject: [PATCH 2/3] fix(gguf): refuse reap-style per-layer expert prunes in deepseek4 ggufs --- python/freetoken/models/deepseek_v4/gguf.py | 36 ++++++++++++++++++--- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py index 50457fcd86..4368939bd6 100644 --- a/python/freetoken/models/deepseek_v4/gguf.py +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -73,7 +73,9 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: def _compress_ratios(index: dict[str, GgufTensorInfo], head_dim: int, num_layers: int) -> tuple[int, ...]: - """Per-layer compressor ratios from the tensor table: the KV placeholder carries a single 0. + """Per-layer compressor ratios from the tensor table, which every released file's + ``attention.compress_ratios`` array (43 trunk entries ahead of 3 MTP zeros, which + llama.cpp requires) confirms element for element. An ``indexer.proj`` tensor exists exactly on the ratio-4 blocks; a compressor without an indexer is a 128 block, unless its ``wkv`` has the CSA overlap's double rows (ratio 4).""" @@ -134,6 +136,22 @@ def need(key: str): num_layers = int(need("block_count")) index = gguf_tensor_index(shim.model_path) + + # REAP exports prune experts per layer: as a per-layer expert_count array, or with the + # scalar kept but the expert tensors shrunk. llama.cpp builds every layer from one + # n_expert and refuses both shapes; the model has no per-layer expert count to express. + experts_kv = need("expert_count") + if isinstance(experts_kv, list): + raise NotImplementedError( + f"{shim.model_path}: a per-layer expert_count array (REAP-style prune) is not served" + ) + n_experts = int(experts_kv) + for name, info in index.items(): + if name.endswith(".ffn_gate_exps.weight") and info.ne[-1] != n_experts: + raise NotImplementedError( + f"{shim.model_path}: {name} carries {info.ne[-1]} experts but expert_count says " + f"{n_experts} (REAP-style per-layer prunes are not served)" + ) args = DeepseekV4Args( max_batch_size=1, vocab_size=int(shim.vocab_size), @@ -181,22 +199,30 @@ def need(key: str): def _float(info: GgufTensorInfo) -> torch.Tensor: """A float tensor in torch order, dequantizing packed ggml rows: dynamic quants quantize - tensors the module consumes as floats (the mHC mixers, the compressor gates, the router).""" + tensors the module consumes as floats (the mHC mixers, the compressor gates, the router). + Trailing unit dims in ``ne`` are dropped -- some writers ship tensors 4-D (llama.cpp's + loaders try that shape first) and a trailing 1 carries no elements.""" + ne = info.ne + while len(ne) > 1 and ne[-1] == 1: + ne = ne[:-1] if not is_quantized(info.ggml_type): from freetoken.models.gguf.tensors import float_tensor - return float_tensor(info) + return float_tensor(info).view(*reversed(ne)) from freetoken.third_party.gguf import quants from freetoken.third_party.gguf.constants import GGMLQuantizationType flat = quants.dequantize(np.array(info.data()), GGMLQuantizationType(info.ggml_type)) - return torch.from_numpy(flat).view(*reversed(info.ne)) + return torch.from_numpy(flat).view(*reversed(ne)) def _tid2eid(info: GgufTensorInfo) -> torch.Tensor: """The hash-route table as the module's int64 [vocab, top_k].""" + ne = info.ne + while len(ne) > 1 and ne[-1] == 1: + ne = ne[:-1] raw = np.frombuffer(info.data(), dtype=np.int32).copy() - return torch.from_numpy(raw).view(*reversed(info.ne)).to(torch.int64) + return torch.from_numpy(raw).view(*reversed(ne)).to(torch.int64) def iter_gguf_weights( From 39f6b7e5d0fe1e3807394dff97fb6128f5dfffb8 Mon Sep 17 00:00:00 2001 From: Mukkesh Ganesh Date: Wed, 21 Oct 2026 09:36:10 -0700 Subject: [PATCH 3/3] fix(gguf): apply the deepseek4 review fixes - route the offload experts through swiglu_clamp so the GGUF MoE path applies the limit - map model.embed -> embed on the HF row so the bf16 exemption matches - default a wholly missing rope.scaling block to the family yarn values with a warning; refuse a partly present one --- python/freetoken/models/deepseek_v4/gguf.py | 32 +++++++++++++++++---- python/freetoken/models/deepseek_v4/moe.py | 4 ++- python/freetoken/models/register.py | 2 +- 3 files changed, 31 insertions(+), 7 deletions(-) diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py index 4368939bd6..4f90137d8b 100644 --- a/python/freetoken/models/deepseek_v4/gguf.py +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -27,6 +27,7 @@ from freetoken.layers.quantization.ggml_types import is_quantized from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index, gguf_tensor_names from freetoken.models.gguf.tensors import check_scales, fused, matrix +from freetoken.utils import init_logger from .args import DeepseekV4Args from .config import model_config_from_args @@ -35,6 +36,8 @@ from freetoken.models.config import ModelConfig from freetoken.models.gguf.config import GgufConfigShim +logger = init_logger(__name__) + # engine module leaf (after "model.layers.N.") -> GGUF tensor names (without ".weight") _LAYER_LEAVES: dict[str, tuple[str, ...]] = { "attn.wq_a": ("attn_q_a",), @@ -125,8 +128,27 @@ def need(key: str): if int(need("expert_gating_func")) != int(ExpertGatingFuncType.SQRTSOFTPLUS): raise NotImplementedError(f"{shim.model_path}: expert_gating_func other than sqrtsoftplus is not served") - if kv("rope.scaling.type", "yarn") != "yarn": + # every official DeepSeek-V4 config carries the same yarn block (V4-Flash, 0731, Pro, + # Pro-0813: factor 16 over 65536, betas 32 / 1). The no-yarn files on HF lost the whole + # block in conversion -- llama.cpp serves them unscaled, which is wrong for them -- so a + # wholly missing block falls back to the family values and a partly present one refuses. + scaling = {k: kv(f"rope.scaling.{k}") for k in + ("type", "factor", "original_context_length", "yarn_beta_fast", "yarn_beta_slow")} + if scaling["type"] not in (None, "yarn"): raise NotImplementedError(f"{shim.model_path}: rope scaling other than yarn is not served") + missing = [k for k, v in scaling.items() if v is None] + if missing: + if len(missing) != len(scaling): + raise ValueError( + f"{shim.model_path}: rope.scaling is only partly present " + f"(missing {', '.join(missing)}); the file is inconsistent" + ) + logger.warning( + f"{shim.model_path}: no rope.scaling block; using the family values " + "(yarn, factor 16, original 65536, betas 32 / 1)" + ) + scaling = {"type": "yarn", "factor": 16.0, "original_context_length": 65536, + "yarn_beta_fast": 32.0, "yarn_beta_slow": 1.0} head_dim = int(need("attention.key_length")) if int(need("attention.value_length")) != head_dim: @@ -175,11 +197,11 @@ def need(key: str): window_size=int(need("attention.sliding_window")), compress_ratios=_compress_ratios(index, head_dim, num_layers), compress_rope_theta=float(kv("attention.compress_rope_freq_base", 160000.0) or 160000.0), - original_seq_len=int(need("rope.scaling.original_context_length")), + original_seq_len=int(scaling["original_context_length"]), rope_theta=float(need("rope.freq_base")), - rope_factor=float(need("rope.scaling.factor")), - beta_fast=float(kv("rope.scaling.yarn_beta_fast", 32) or 32), - beta_slow=float(kv("rope.scaling.yarn_beta_slow", 1) or 1), + rope_factor=float(scaling["factor"]), + beta_fast=float(scaling["yarn_beta_fast"]), + beta_slow=float(scaling["yarn_beta_slow"]), index_n_heads=int(need("attention.indexer.head_count")), index_head_dim=int(need("attention.indexer.key_length")), index_topk=int(need("attention.indexer.top_k")), diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index aadcdc0481..4c58844c88 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -97,7 +97,9 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, *, strategy: str = "offl hidden_size=args.dim, intermediate_size=args.moe_inter_dim, renormalize=True, - activation="silu", + # the clamped SwiGLU the reference runs (up to +-limit, gate to max limit); the + # GGUF MoE path keys on this name and plain "silu" would drop the clamp + activation="swiglu_clamp", limit=args.swiglu_limit, strategy=strategy, decode_target=decode_target, diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 65601d63db..f4238668a1 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -206,7 +206,7 @@ class ModelSpec: "freetoken.models.deepseek_v4", "DeepseekV4ForCausalLM", # the checkpoint has no ``model.`` root - checkpoint_roots=(("model.layers", "layers"), ("model.head", "head")), + checkpoint_roots=(("model.layers", "layers"), ("model.embed", "embed"), ("model.head", "head")), packed_modules_mapping=_EXPERTS_W123_PACKED, # the embed, the head, the KV compressors, the indexer's scorer and the vision tower + # aligner ship bf16; the fp8 config has no modules_to_not_convert (patterns match