From 85ea907fc27be41e88b9fcd34456566ca40ab3c1 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Fri, 16 Oct 2026 08:49:13 +0000 Subject: [PATCH] feat(gguf): serve Qwen2, Qwen3 and Qwen3 MoE GGUF checkpoints --- python/freetoken/kernel/aot_models.py | 9 +- python/freetoken/models/gguf/config.py | 3 + python/freetoken/models/qwen2/__init__.py | 3 +- python/freetoken/models/qwen2/gguf.py | 144 +++++++++++++++++ python/freetoken/models/qwen2/model.py | 2 + python/freetoken/models/qwen3/__init__.py | 3 +- python/freetoken/models/qwen3/gguf.py | 143 +++++++++++++++++ python/freetoken/models/qwen3/model.py | 2 + python/freetoken/models/qwen3_moe/__init__.py | 4 + python/freetoken/models/qwen3_moe/gguf.py | 149 ++++++++++++++++++ python/freetoken/models/qwen3_moe/model.py | 2 + python/freetoken/models/register.py | 22 +++ 12 files changed, 480 insertions(+), 6 deletions(-) create mode 100644 python/freetoken/models/qwen2/gguf.py create mode 100644 python/freetoken/models/qwen3/gguf.py create mode 100644 python/freetoken/models/qwen3_moe/gguf.py diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index 308c2d6810..8c5d32ba9d 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -142,7 +142,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int moe_intermediate_size=768, expert_formats=("bf16",), aliases=("Qwen/Qwen3-30B-A3B-Thinking-2507",), - arch_aliases=("Qwen3VLMoeForConditionalGeneration",), # Qwen3-VL-30B-A3B: same text tower + arch_aliases=("Qwen3VLMoeForConditionalGeneration", "Qwen3MoeGGUFForCausalLM"), # Qwen3-VL-30B-A3B: same text tower ), AotModel( name="Qwen/Qwen3.5-35B-A3B", @@ -152,7 +152,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int top_k=8, moe_intermediate_size=512, expert_formats=("bf16",), - arch_aliases=("Qwen3_5MoeForCausalLM",), # text-only release of the same tower + arch_aliases=("Qwen3_5MoeForCausalLM", "Qwen3_5MoeGGUFForCausalLM"), # text-only release of the same tower ), AotModel( name="Qwen/Qwen3.5-35B-A3B-FP8", @@ -348,7 +348,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int hidden_size=5120, kv_groups=((4, 256),), aliases=("Qwen/Qwen3.6-27B-FP8", "nvidia/Qwen3.6-27B-NVFP4"), - arch_aliases=("Qwen3_5ForCausalLM",), # text-only release of the same tower + arch_aliases=("Qwen3_5ForCausalLM", "Qwen3_5GGUFForCausalLM"), # text-only release of the same tower ), AotModel( name="google/gemma-4-12B-it", @@ -381,13 +381,14 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int architecture="Qwen2ForCausalLM", hidden_size=3584, kv_groups=((4, 128),), + arch_aliases=("Qwen2GGUFForCausalLM",), ), AotModel( name="Qwen/Qwen3-8B", architecture="Qwen3ForCausalLM", hidden_size=4096, kv_groups=((8, 128),), - arch_aliases=("Qwen3VLForConditionalGeneration",), # Qwen3-VL-8B: same text tower + arch_aliases=("Qwen3VLForConditionalGeneration", "Qwen3GGUFForCausalLM"), # Qwen3-VL-8B: same text tower ), AotModel( name="mistralai/Mistral-7B-Instruct-v0.3", diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index f9210938fb..518acfe4a6 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -18,6 +18,9 @@ # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", + "qwen2": "Qwen2GGUFForCausalLM", + "qwen3": "Qwen3GGUFForCausalLM", + "qwen3moe": "Qwen3MoeGGUFForCausalLM", "qwen35": "Qwen3_5GGUFForCausalLM", "qwen35moe": "Qwen3_5MoeGGUFForCausalLM", "qwen4exp": "Qwen4ExpGGUFForCausalLM", diff --git a/python/freetoken/models/qwen2/__init__.py b/python/freetoken/models/qwen2/__init__.py index 43ca84f82a..3917c188fd 100644 --- a/python/freetoken/models/qwen2/__init__.py +++ b/python/freetoken/models/qwen2/__init__.py @@ -1,5 +1,6 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Qwen2ForCausalLM from .weight import iter_weights -__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights"] +__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen2/gguf.py b/python/freetoken/models/qwen2/gguf.py new file mode 100644 index 0000000000..501c113a3a --- /dev/null +++ b/python/freetoken/models/qwen2/gguf.py @@ -0,0 +1,144 @@ +"""Qwen2 / Qwen2.5 from a llama.cpp ``qwen2`` GGUF: config from the GGUF metadata, tensors by name. + +The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen2" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint 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 () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"} + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling == "yarn": + rope.update( + rope_type="yarn", + factor=float(kv("rope.scaling.factor")), + original_max_position_embeddings=int(kv("rope.scaling.original_context_length")), + ) + elif scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + # Qwen2.5 files carry no key_length; parse_config then derives head_dim from the hidden size + key_length = md.get(f"{_ARCH}.attention.key_length") + return SimpleNamespace( + model_type="qwen2", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=int(kv("block_count")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(key_length) if key_length is not None else None, + rope_scaling=rope, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return parse_config(_hf_config(shim)) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The state dict from the GGUF.""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = hf.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + bias = index.get(f"{blk}.attn_qkv.bias") + yield f"{sa}.qkv_proj.bias", float_tensor(bias) if bias is not None else torch.cat([float_tensor(t(f"{blk}.attn_{p}.bias")) for p in ("q", "k", "v")]) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + mlp = f"{pre}.mlp" + gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0])) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen2/model.py b/python/freetoken/models/qwen2/model.py index ed56fde25a..94bc2beb6b 100644 --- a/python/freetoken/models/qwen2/model.py +++ b/python/freetoken/models/qwen2/model.py @@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/qwen3/__init__.py b/python/freetoken/models/qwen3/__init__.py index 71bda9781f..fef7c8fb99 100644 --- a/python/freetoken/models/qwen3/__init__.py +++ b/python/freetoken/models/qwen3/__init__.py @@ -1,5 +1,6 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Qwen3ForCausalLM from .weight import iter_weights -__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights"] +__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3/gguf.py b/python/freetoken/models/qwen3/gguf.py new file mode 100644 index 0000000000..06245ce4c2 --- /dev/null +++ b/python/freetoken/models/qwen3/gguf.py @@ -0,0 +1,143 @@ +"""Qwen3 dense from a llama.cpp ``qwen3`` GGUF: config from the GGUF metadata, tensors by name. + +The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen3" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate_proj": "ffn_gate", + "mlp.up_proj": "ffn_up", + "mlp.down_proj": "ffn_down", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint 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 () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"} + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling == "yarn": + rope.update( + rope_type="yarn", + factor=float(kv("rope.scaling.factor")), + original_max_position_embeddings=int(kv("rope.scaling.original_context_length")), + ) + elif scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + key_length = md.get(f"{_ARCH}.attention.key_length") + return SimpleNamespace( + model_type="qwen3", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=int(kv("block_count")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(key_length) if key_length is not None else None, + rope_scaling=rope, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return parse_config(_hf_config(shim)) + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The state dict from the GGUF.""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = hf.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm")) + yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm")) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + mlp = f"{pre}.mlp" + gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden) + yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0])) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3/model.py b/python/freetoken/models/qwen3/model.py index 6ac7d2fa8f..ae0d4394f5 100644 --- a/python/freetoken/models/qwen3/model.py +++ b/python/freetoken/models/qwen3/model.py @@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/qwen3_moe/__init__.py b/python/freetoken/models/qwen3_moe/__init__.py index 59c74fb0cb..cb96a1921b 100644 --- a/python/freetoken/models/qwen3_moe/__init__.py +++ b/python/freetoken/models/qwen3_moe/__init__.py @@ -1,4 +1,5 @@ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import Qwen3MoeForCausalLM from .weight import iter_weights, iter_weights_parallel @@ -7,4 +8,7 @@ "parse_config", "iter_weights", "iter_weights_parallel", + "iter_gguf_weights", + "parse_gguf_config", + "to_gguf", ] diff --git a/python/freetoken/models/qwen3_moe/gguf.py b/python/freetoken/models/qwen3_moe/gguf.py new file mode 100644 index 0000000000..7bf016292e --- /dev/null +++ b/python/freetoken/models/qwen3_moe/gguf.py @@ -0,0 +1,149 @@ +"""Qwen3 MoE from a llama.cpp ``qwen3moe`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``). + +The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones. +""" + +from __future__ import annotations + +import re +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import float_tensor, fused, matrix + +from .config import parse_config + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen3moe" +_LAYER_LEAVES = { + "self_attn.q_proj": "attn_q", + "self_attn.k_proj": "attn_k", + "self_attn.v_proj": "attn_v", + "self_attn.o_proj": "attn_output", + "mlp.gate": "ffn_gate_inp", + "mlp.experts.0.gate_proj": "ffn_gate_exps", + "mlp.experts.0.up_proj": "ffn_up_exps", + "mlp.experts.0.down_proj": "ffn_down_exps", +} +_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj") +_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Checkpoint 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 () + blk, leaf = f"blk.{m.group(1)}", m.group(2) + if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors: + return (f"{blk}.attn_qkv",) + if leaf in ("mlp.experts.0.gate_proj", "mlp.experts.0.up_proj") and f"{blk}.ffn_gate_up_exps.weight" in tensors: + return (f"{blk}.ffn_gate_up_exps",) + tensor = _LAYER_LEAVES.get(leaf) + return (f"{blk}.{tensor}",) if tensor is not None else () + + +def _check_supported(shim: GgufConfigShim) -> None: + if any(k.startswith("prism.hadamard.") for k in shim.metadata): + raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported") + scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale"))) + if scaled: + raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})") + + +def _hf_config(shim: GgufConfigShim) -> SimpleNamespace: + """The HF config fields ``parse_config`` reads, from the GGUF metadata.""" + md = shim.metadata + + def kv(key: str): + return md[f"{_ARCH}.{key}"] + + _check_supported(shim) + rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"} + scaling = md.get(f"{_ARCH}.rope.scaling.type", "none") + if scaling == "yarn": + rope.update( + rope_type="yarn", + factor=float(kv("rope.scaling.factor")), + original_max_position_embeddings=int(kv("rope.scaling.original_context_length")), + ) + elif scaling != "none": + raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported") + key_length = md.get(f"{_ARCH}.attention.key_length") + return SimpleNamespace( + model_type="qwen3_moe", + architectures=list(shim.architectures), + hidden_act="silu", + hidden_size=int(kv("embedding_length")), + num_hidden_layers=int(kv("block_count")), + num_attention_heads=int(kv("attention.head_count")), + num_key_value_heads=int(kv("attention.head_count_kv")), + head_dim=int(key_length) if key_length is not None else None, + rope_scaling=rope, + max_position_embeddings=int(kv("context_length")), + rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")), + vocab_size=int(shim.vocab_size), + tie_word_embeddings=bool(shim.tie_word_embeddings), + intermediate_size=int(kv("feed_forward_length")), + num_experts=int(kv("expert_count")), + num_experts_per_tok=int(kv("expert_used_count")), + moe_intermediate_size=int(kv("expert_feed_forward_length")), + # llama.cpp always renormalizes the top-k weights of a qwen3moe router + norm_topk_prob=True, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + return replace(parse_config(_hf_config(shim)), expert_quant="gguf") + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).""" + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + hf = _hf_config(cached_load_hf_config(model_path)) + index = gguf_tensor_index(model_path) + hidden = hf.hidden_size + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden) + if not hf.tie_word_embeddings: + yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden) + yield "model.norm.weight", float_tensor(t("output_norm")) + + for layer in range(hf.num_hidden_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm")) + yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm")) + sa = f"{pre}.self_attn" + # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix + qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden) + yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm")) + yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm")) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + yield f"{pre}.mlp.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + + +__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3_moe/model.py b/python/freetoken/models/qwen3_moe/model.py index 5d89656bda..7d81dd2673 100644 --- a/python/freetoken/models/qwen3_moe/model.py +++ b/python/freetoken/models/qwen3_moe/model.py @@ -47,6 +47,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"): self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, + quant_config=config.quant, + prefix=f"{prefix}.embed_tokens", ) self.layers = OPList( [ diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 86e3a734c8..263b5f932c 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -110,6 +110,28 @@ class ModelSpec: "Qwen3MoeForCausalLM", packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED, ), + # llama.cpp qwen2 / qwen3 / qwen3moe GGUF: the same model classes with every tensor in its packed ggml type + "Qwen2GGUFForCausalLM": ModelSpec( + "freetoken.models.qwen2", + "Qwen2ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED, + ), + "Qwen3GGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3", + "Qwen3ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED, + ), + "Qwen3MoeGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_moe", + "Qwen3MoeForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + packed_modules_mapping=_DENSE_PACKED + _EXPERTS_PACKED, + ), # Qwen3-VL: the Qwen3 text tower under model.language_model. plus the shared Qwen VL vision tower with DeepStack; the MoE variant ships its experts pre-stacked. "Qwen3VLForConditionalGeneration": ModelSpec( "freetoken.models.qwen3_vl",