diff --git a/docs/cli.md b/docs/cli.md index 350c03d45..94abe5042 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -102,11 +102,11 @@ See [models.md](models.md#moe-strategies) for what each strategy does, and [perf | Flag | Default | Meaning | |---|---|---| -| `--speculative-algorithm`, `--spec-algo` | off | Draft method. `mtp` drafts with the checkpoint's multi-token-prediction layer: Qwen3.5 / Qwen3.6 HF checkpoints, and Qwen3.8-Flash-Next HF or GGUF | +| `--speculative-algorithm`, `--spec-algo` | off | Draft method. `mtp` drafts with the checkpoint's multi-token-prediction layer: Qwen3.5 / Qwen3.6 HF or GGUF checkpoints, and Qwen3.8-Flash-Next HF or GGUF | | `--speculative-num-steps`, `--spec-num-steps` | the method's default (3 for `mtp`) | Draft tokens per decode step, 1 to 8 for `mtp` | -| `--speculative-draft-model-path`, `--spec-draft-model` | the target checkpoint | Where the draft weights load from. A Qwen3.8-Flash-Next target also takes a Strata MTP pack (a GGUF that keeps the checkpoint's `mtp.*` tensors under their HF names), and a GGUF target needs one; the `nextn` MTP blocks llama.cpp writes are not read; a separate checkpoint folder is read with the target's quantization config | +| `--speculative-draft-model-path`, `--spec-draft-model` | the target checkpoint | Where the draft weights load from. A Qwen3.8-Flash-Next target also takes a Strata MTP pack (a GGUF that keeps the checkpoint's `mtp.*` tensors under their HF names), and a Qwen3.8-Flash-Next GGUF target needs one; a Qwen3.5 / Qwen3.6 GGUF target drafts from its own `nextn` MTP block (the block llama.cpp keeps after the last decoder layer), or from an `mtp-*.gguf` sidecar of the same model (unsloth's `MTP/` releases) given here; a separate checkpoint folder is an HF-target option, read with the target's quantization config, and a GGUF target refuses a folder draft | -Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets with sliding-window or DSV4 window-tier KV are refused at startup. `mtp` also refuses Qwen3.5 / Qwen3.6 GGUF targets and an FTW checkpoint as the draft source, since FTW drops the `mtp.*` tensors. +Speculative decoding needs `--tp 1` and runs without overlap scheduling; targets with sliding-window or DSV4 window-tier KV are refused at startup. `mtp` refuses an FTW checkpoint as the draft source, since FTW drops the `mtp.*` tensors, and a Qwen3.5 / Qwen3.6 GGUF converted without its MTP layer (llama.cpp `--no-mtp`) when no `mtp-*.gguf` sidecar is given: such a trunk has no draft of its own to read, while a `--no-mtp` trunk plus a sidecar works (the trunk config's shapes, the sidecar's weights). ### API behaviour diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 2b80e57f1..6f52d8885 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from freetoken.models.config import ModelConfig + from freetoken.moe.expert_banks import ExpertBanks from freetoken.models.gguf.config import GgufConfigShim _LAYER_LEAVES: dict[str, tuple[str, ...]] = { @@ -51,19 +52,37 @@ "lm_head": ("output",), } _LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.(.+)$") +# the MTP draft's layer reads the block after the trunk: the same leaves under mtp.layers.0 +_MTP_LAYER_RE = re.compile(r"^mtp\.layers\.0\.(.+)$") + + +def _block_tensors(blk: str, leaf: str, tensors: dict[str, int]) -> tuple[str, ...]: + """The GGUF tensors (without ``.weight``) of one block's module leaf; stacked experts fuse gate|up when the file stores them fused.""" + 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",) + return tuple(f"{blk}.{t}" for t in _LAYER_LEAVES.get(leaf, ())) + + +def _nextn_blk(tensors: dict[str, int]) -> str | None: + """The block whose ``nextn`` tensors the MTP draft reads (``blk.64``), or None in a file converted without them.""" + return next((name.split(".nextn.")[0] for name in tensors if ".nextn.eh_proj.weight" in name), None) 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] + if name == "mtp.fc": + nextn = _nextn_blk(tensors) + return () if nextn is None else (f"{nextn}.nextn.eh_proj",) + m = _MTP_LAYER_RE.match(name) + if m is not None: + nextn = _nextn_blk(tensors) + return () if nextn is None else _block_tensors(nextn, m.group(1), tensors) m = _LAYER_RE.match(name) if m is None: return () - blk, leaf = f"blk.{m.group(1)}", m.group(2) - 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",) - return tuple(f"{blk}.{t}" for t in _LAYER_LEAVES.get(leaf, ())) + return _block_tensors(f"blk.{m.group(1)}", m.group(2), tensors) def _check_supported(shim: GgufConfigShim) -> None: @@ -144,6 +163,63 @@ def _untile(t: torch.Tensor, num_k_heads: int, head_dim: int, dim: int = 0) -> t return t.index_select(dim, idx) +def _iter_block( + t, + norm, + blk: str, + pre: str, + config: ModelConfig, + layer_id: int, +) -> Iterator[tuple[str, torch.Tensor]]: + """One decoder block (attention, MLP) as ``{pre}.*`` state-dict entries; the trunk loop and the nextn draft reader share it.""" + hidden = config.hidden_size + yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm") + yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.post_attention_norm") + if config.is_linear_layer(layer_id): + g = config.linear_attention_group() + la = f"{pre}.linear_attn" + nk, dk = g.num_key_heads, g.key_head_dim + qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate") + qkv_rows = matrix(qkv) + v_start = 2 * nk * dk + qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)]) + z_rows = _untile(matrix(gate), nk, g.value_head_dim) + beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha") + beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info)) + if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type): + yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden) + yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden) + else: + yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden) + conv = float_tensor(t(f"{blk}.ssm_conv1d")).float() + conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)]) + yield f"{la}.conv1d.weight", conv.unsqueeze(1) + yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1) + yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1) + yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm")) + yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim) + else: + sa = f"{pre}.self_attn" + parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] + yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden) + out = t(f"{blk}.attn_output") + yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) + yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm") + yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm") + + mlp = f"{pre}.mlp" + if config.moe_enabled: + yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) + sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden) + yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size) + yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1) + else: + fg, fu, fd = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) + yield from fused(f"{mlp}.gate_up_proj", [fg, fu], [matrix(fg), matrix(fu)], hidden) + yield from fused(f"{mlp}.down_proj", [fd], [matrix(fd)], config.intermediate_size) + + def iter_gguf_weights( model_path: str, device: torch.device, @@ -161,7 +237,6 @@ def iter_gguf_weights( vision = include_vision and config.is_multimodal log_mmproj(model_path, vision) index = gguf_tensor_index(model_path) - g = config.linear_attention_group() hidden = config.hidden_size def t(name: str) -> GgufTensorInfo: @@ -176,54 +251,77 @@ def norm(name: str) -> torch.Tensor: yield "model.norm.weight", norm("output_norm") for layer in range(config.num_layers): - blk, pre = f"blk.{layer}", f"model.layers.{layer}" - yield f"{pre}.input_layernorm.weight", norm(f"{blk}.attn_norm") - yield f"{pre}.post_attention_layernorm.weight", norm(f"{blk}.post_attention_norm") - if config.is_linear_layer(layer): - la = f"{pre}.linear_attn" - nk, dk = g.num_key_heads, g.key_head_dim - qkv, gate = t(f"{blk}.attn_qkv"), t(f"{blk}.attn_gate") - qkv_rows = matrix(qkv) - v_start = 2 * nk * dk - qkv_rows = torch.cat([qkv_rows[:v_start], _untile(qkv_rows[v_start:], nk, g.value_head_dim)]) - z_rows = _untile(matrix(gate), nk, g.value_head_dim) - beta_info, alpha_info = t(f"{blk}.ssm_beta"), t(f"{blk}.ssm_alpha") - beta, alpha = (_untile(matrix(i), nk, 1) for i in (beta_info, alpha_info)) - if is_quantized(qkv.ggml_type) or is_quantized(gate.ggml_type): - yield from fused(f"{la}.in_proj_qkvz", [qkv, gate], [qkv_rows, z_rows], hidden) - yield from fused(f"{la}.in_proj_ba", [beta_info, alpha_info], [beta, alpha], hidden) - else: - yield from fused(f"{la}.in_proj", [qkv, gate, beta_info, alpha_info], [qkv_rows, z_rows, beta, alpha], hidden) - conv = float_tensor(t(f"{blk}.ssm_conv1d")).float() - conv = torch.cat([conv[:v_start], _untile(conv[v_start:], nk, g.value_head_dim)]) - yield f"{la}.conv1d.weight", conv.unsqueeze(1) - yield f"{la}.dt_bias", _untile(float_tensor(t(f"{blk}.ssm_dt.bias")).float(), nk, 1) - yield f"{la}.A_log", _untile(torch.log(-float_tensor(t(f"{blk}.ssm_a")).float()), nk, 1) - yield f"{la}.norm.weight", float_tensor(t(f"{blk}.ssm_norm")) - yield from fused(f"{la}.out_proj", [t(f"{blk}.ssm_out")], [matrix(t(f"{blk}.ssm_out"))], g.num_value_heads * g.value_head_dim) - else: - sa = f"{pre}.self_attn" - parts = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")] - yield from fused(f"{sa}.qkv_proj", parts, [matrix(p) for p in parts], hidden) - out = t(f"{blk}.attn_output") - yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0])) - yield f"{sa}.q_norm.weight", norm(f"{blk}.attn_q_norm") - yield f"{sa}.k_norm.weight", norm(f"{blk}.attn_k_norm") - - mlp = f"{pre}.mlp" - if config.moe_enabled: - yield f"{mlp}.gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp")) - sg, su, sd = (t(f"{blk}.ffn_{p}_shexp") for p in ("gate", "up", "down")) - yield from fused(f"{mlp}.shared_expert.gate_up_proj", [sg, su], [matrix(sg), matrix(su)], hidden) - yield from fused(f"{mlp}.shared_expert.down_proj", [sd], [matrix(sd)], config.shared_expert_intermediate_size) - yield f"{mlp}.shared_expert_gate.weight", float_tensor(t(f"{blk}.ffn_gate_inp_shexp")).reshape(1, -1) - else: - fg, fu, fd = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down")) - yield from fused(f"{mlp}.gate_up_proj", [fg, fu], [matrix(fg), matrix(fu)], hidden) - yield from fused(f"{mlp}.down_proj", [fd], [matrix(fd)], config.intermediate_size) + yield from _iter_block(t, norm, f"blk.{layer}", f"model.layers.{layer}", config, layer) if vision: yield from iter_mmproj_weights(find_mmproj(model_path)) -__all__ = ["gguf_mm_config", "iter_gguf_weights", "parse_gguf_config", "to_gguf"] +# ---------------------------------------------------------------------------------------------- +# The trunk's nextn (MTP) block: llama.cpp keeps it after the last decoder layer +# ---------------------------------------------------------------------------------------------- + +def nextn_block_index(model_path: str) -> int: + """The trunk block that carries the nextn tensors; parse_gguf_config subtracts the same count from block_count.""" + from freetoken.models.gguf.reader import gguf_architecture, load_gguf_metadata + + md = load_gguf_metadata(model_path) + arch = gguf_architecture(model_path) + nextn = int(md.get(f"{arch}.nextn_predict_layers", 0)) + if nextn == 0: + raise ValueError( + f"{model_path}: no MTP (nextn) block to draft from: the trunk was converted without one " + "(llama.cpp --no-mtp), or the file is a Strata-style MTP pack, which is not a draft " + "source for this family" + ) + if nextn > 1: + raise ValueError(f"{model_path}: {nextn} nextn layers, but the MTP draft reads exactly one") + return int(md[f"{arch}.block_count"]) - nextn + + +def iter_nextn_weights(model_path: str, config: ModelConfig) -> Iterator[tuple[str, torch.Tensor]]: + """The nextn block as Qwen3_5MTPDraft names it: the block's tensors under the draft's state-dict names, its routed experts excluded (they load as resident banks).""" + blk = f"blk.{nextn_block_index(model_path)}" + index = gguf_tensor_index(model_path) + if f"{blk}.nextn.eh_proj.weight" not in index: + raise ValueError(f"{model_path}: {blk} has no nextn tensors") + + def t(name: str) -> GgufTensorInfo: + key = name if name in index else f"{name}.weight" + if key not in index: + raise ValueError(f"{model_path}: the MTP (nextn) block lacks {key}; reconvert with a current llama.cpp") + return index[key] + + def norm(name: str) -> torch.Tensor: + return float_tensor(t(name)).float() - 1.0 + + eh = t(f"{blk}.nextn.eh_proj") + yield from fused("fc", [eh], [matrix(eh)], 2 * config.hidden_size) + yield "pre_fc_norm_embedding.weight", norm(f"{blk}.nextn.enorm") + yield "pre_fc_norm_hidden.weight", norm(f"{blk}.nextn.hnorm") + yield "norm.weight", norm(f"{blk}.nextn.shared_head_norm") + # the draft's one layer is full attention, whatever the trunk mixes around it + yield from _iter_block(t, norm, blk, "layers.0", config, 0) + + +def nextn_expert_banks(model_path: str, device: torch.device) -> "ExpertBanks": + """The nextn block's stacked experts as the draft's one resident ggml bank set.""" + from freetoken.models.gguf.experts import load_gguf_expert_banks + + return load_gguf_expert_banks(model_path, 1, first_layer=nextn_block_index(model_path), device=device) + + +def nextn_quant(model_path: str) -> "GGUFConfig": + """The draft's quant over its own file's tensor table: the trunk's when the draft is the + trunk's nextn block, the sidecar's when it is an mtp-*.gguf at another ggml mix.""" + from freetoken.models.gguf.config import GGUF_ARCH_TO_REGISTRY + from freetoken.models.gguf.reader import gguf_architecture + from freetoken.models.register import checkpoint_quant_config, get_model_spec + + # the HF and GGUF rows share one model class, so the target's class name resolves to the + # HF row and its from_hf branch; the file's own architecture names the GGUF row instead + spec = get_model_spec(GGUF_ARCH_TO_REGISTRY[gguf_architecture(model_path)]) + return checkpoint_quant_config(model_path, SimpleNamespace(model_path=model_path), spec) + + +__all__ = ["gguf_mm_config", "iter_gguf_weights", "iter_nextn_weights", "nextn_block_index", "nextn_expert_banks", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/qwen3_5_moe/mtp.py b/python/freetoken/models/qwen3_5_moe/mtp.py index d3b257b65..20ae6c52a 100644 --- a/python/freetoken/models/qwen3_5_moe/mtp.py +++ b/python/freetoken/models/qwen3_5_moe/mtp.py @@ -46,12 +46,41 @@ def relay_width(config: ModelConfig) -> int: @staticmethod def check_source(path: str, target_path: str) -> None: - from freetoken.models.gguf.reader import is_gguf_path + from freetoken.models.gguf.reader import gguf_architecture, gguf_tensor_index, is_gguf_path, load_gguf_metadata from freetoken.models.loader import safetensors_weight_map if is_gguf_path(path): - raise ValueError(f"the Qwen3.5 MTP draft loads from safetensors, not from the GGUF {path}") - if os.path.isdir(path) and not any(name.startswith("mtp.") for name in safetensors_weight_map(path)): + if not is_gguf_path(target_path): + raise ValueError(f"the Qwen3.5 MTP draft loads from safetensors mtp.* tensors, not from the GGUF {path}") + from .gguf import nextn_block_index + + blk = nextn_block_index(path) + index = gguf_tensor_index(path) + if f"blk.{blk}.nextn.eh_proj.weight" not in index: + raise ValueError(f"{path}: blk.{blk} carries no nextn tensors") + if os.path.abspath(path) != os.path.abspath(target_path): + # a separate draft file must be an mtp-*.gguf sidecar: only the MTP block, same + # arch; anything else is another trunk read under this target's shapes + blocks = {name.split(".")[1] for name in index if name.startswith("blk.")} + if blocks != {str(blk)} or gguf_architecture(path) != gguf_architecture(target_path): + raise ValueError( + f"the MTP draft of a GGUF target is the trunk's own nextn block or an " + f"mtp-*.gguf sidecar of the same model, not {path}" + ) + # the same arch string spans sizes (qwen35 covers the 9B and the 27B); a wrong-size + # sidecar would otherwise die on a bare shape assert deep in load_state_dict + side_hidden = load_gguf_metadata(path).get(f"{gguf_architecture(path)}.embedding_length") + target_hidden = load_gguf_metadata(target_path).get(f"{gguf_architecture(target_path)}.embedding_length") + if side_hidden is not None and target_hidden is not None and int(side_hidden) != int(target_hidden): + raise ValueError( + f"{path}: an mtp-*.gguf sidecar of a model with hidden size {int(side_hidden)}, but the " + f"target {target_path} has hidden size {int(target_hidden)}; the sidecar must be the " + f"MTP layer of the same model as the target" + ) + elif is_gguf_path(target_path): + # a folder draft would build under the trunk's GGUF quant names + raise ValueError(f"a GGUF target's MTP draft is the trunk's own nextn block, not {path}") + elif os.path.isdir(path) and not any(name.startswith("mtp.") for name in safetensors_weight_map(path)): raise ValueError(f"{path} has no mtp.* tensors") def __init__(self, config: ModelConfig, target_model) -> None: @@ -67,11 +96,26 @@ def __init__(self, config: ModelConfig, target_model) -> None: @classmethod def load(cls, config: ModelConfig, target_model, model_path: str, device: torch.device, dtype: torch.dtype) -> Qwen3_5MTPDraft: + from freetoken.models.gguf.reader import is_gguf_path + + # a GGUF trunk keeps the draft in its own ggml types, mapped by the target config's own to_gguf; + # an mtp-*.gguf sidecar carries a different mix, so its table names the draft's types + gguf = is_gguf_path(model_path) + if gguf: + from .gguf import nextn_quant + + config = dataclasses.replace(config, quant=nextn_quant(model_path)) with torch.device("meta"), torch_dtype(dtype): draft = cls(config, target_model) expected = draft.state_dict() loaded = {} - for name, tensor in iter_mtp_weights(model_path, device): + if gguf: + from .gguf import iter_nextn_weights + + weights = iter_nextn_weights(model_path, draft.config) + else: + weights = iter_mtp_weights(model_path, device) + for name, tensor in weights: if name not in expected: raise ValueError(f"{model_path}: unexpected MTP tensor {name!r}") loaded[name] = tensor.to(device=device, dtype=expected[name].dtype) @@ -83,8 +127,13 @@ def load(cls, config: ModelConfig, target_model, model_path: str, device: torch. # packed by the layer's own expert method, as the target's resident layers are moe = draft.layers.op_list[0].mlp.experts - pieces = iter_mtp_expert_pieces(model_path) - banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True) + if gguf: + from .gguf import nextn_expert_banks + + banks = nextn_expert_banks(model_path, device) + else: + pieces = iter_mtp_expert_pieces(model_path) + banks = build_expert_banks(moe.quant_method, 1, pieces, device=device, resident=True) attach_resident_banks([moe], banks) # not the root: it reaches the target's embedding, which the target finalized already finalize_quant(draft.fc) diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 1fe9cadcd..92722cec9 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -310,7 +310,7 @@ class ModelSpec: spec_drafts=_QWEN3_5_DRAFTS, ), # llama.cpp qwen35 / qwen35moe GGUF: the same model classes with every tensor in its packed ggml type; the vision tower comes from an mmproj file and runs in bf16 - # no MTP draft yet: the nextn blocks and mtp-*.gguf files need a loader of their own + # the MTP draft reads the nextn block llama.cpp keeps after the last decoder layer "Qwen3_5GGUFForCausalLM": ModelSpec( "freetoken.models.qwen3_5_moe", "Qwen3_5ForConditionalGeneration", @@ -321,6 +321,7 @@ class ModelSpec: unquantized_modules=_QWEN3_5_UNQUANTIZED, mm_processor=_QWEN_VL_PROCESSOR, encoders=_QWEN_VL_ENCODERS, + spec_drafts=_QWEN3_5_DRAFTS, ), "Qwen3_5MoeGGUFForCausalLM": ModelSpec( "freetoken.models.qwen3_5_moe", @@ -332,6 +333,7 @@ class ModelSpec: unquantized_modules=_QWEN3_5_UNQUANTIZED, mm_processor=_QWEN_VL_PROCESSOR, encoders=_QWEN_VL_ENCODERS, + spec_drafts=_QWEN3_5_DRAFTS, ), # Qwen3.8-Flash-Next (model_type qwen4_exp): multimodal wrapper config (text tower in # text_config, weights under model.language_model.). 36 GDN + 12 QSA diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py index 381e5a4f4..b57a42fca 100644 --- a/tests/models/test_qwen3_5_moe_weight.py +++ b/tests/models/test_qwen3_5_moe_weight.py @@ -8,6 +8,8 @@ from __future__ import annotations import json +import struct +from types import SimpleNamespace import pytest import torch @@ -19,6 +21,7 @@ from freetoken.models.qwen3_5_moe.config import parse_config from freetoken.models.qwen3_5_moe.weight import iter_expert_pieces, iter_weights, nvfp4_expert_spec from freetoken.models.register import checkpoint_quant_config, get_model_spec +from freetoken.third_party import gguf as gguf_lib from freetoken.utils import cached_load_hf_config H, V = 128, 256 # hidden_size (every block-fp8 projection needs in/out multiples of 128), vocab @@ -664,3 +667,274 @@ def test_mtp_head_fills_the_draft_and_its_experts(checkpoint): assert set(piece) == {"gate", "up", "down", "gate_scale", "up_scale", "down_scale"} assert _same(piece["down"][0], raw[f"{MTP}.mlp.experts.{expert}.down_proj.weight"]) assert _same(piece["gate_scale"][0], raw[f"{MTP}.mlp.experts.{expert}.gate_proj.weight_scale_inv"]) + + +# --------------------------------------------------------------------------- the trunk's nextn block (GGUF) + + +F32, BF16, Q8_0 = 0, 30, int(gguf_lib.GGMLQuantizationType.Q8_0) + + +def _gguf_bytes(tensor: torch.Tensor, ggml_type: int) -> bytes: + if ggml_type == F32: + return tensor.float().numpy().tobytes() + if ggml_type == BF16: + return tensor.to(torch.bfloat16).view(torch.uint16).numpy().tobytes() + raise KeyError(ggml_type) + + +def _write_gguf(path, arch: str, kv: dict[str, object], tensors: dict[str, tuple[int, tuple[int, ...], bytes]]) -> None: + """A minimal GGUF: header, the given KVs (u32 / string) and tensors; the reader needs no more.""" + out = bytearray(b"GGUF" + struct.pack(" None: + key = name.encode() + out.extend(struct.pack(" None: + raw[name] = tensor.float() + tensors[f"blk.{blk}.{name}.weight"] = (ggml_type, tuple(reversed(tensor.shape)), _gguf_bytes(tensor, ggml_type)) + + for name in ("attn_norm", "post_attention_norm", "attn_q_norm", "attn_k_norm"): + put(name, torch.randn({"attn_norm": H, "post_attention_norm": H, "attn_q_norm": AHD, "attn_k_norm": AHD}[name]), F32) + if q_quantized: + raw["attn_q"] = q8 + tensors[f"blk.{blk}.attn_q.weight"] = (Q8_0, (H, q_shape[0]), q8.numpy().tobytes()) + else: + put("attn_q", torch.randn(*q_shape)) + put("attn_k", torch.randn(KVH * AHD, H)) + put("attn_v", torch.randn(KVH * AHD, H)) + put("attn_output", torch.randn(H, QH * AHD)) + if moe: + put("ffn_gate_inp", torch.randn(E, H)) + put("ffn_gate_inp_shexp", torch.randn(H)) + for name, shape in (("ffn_gate_shexp", (I, H)), ("ffn_up_shexp", (I, H)), ("ffn_down_shexp", (H, I))): + put(name, torch.randn(*shape)) + if experts_quantized: + for name, rows, k in (("ffn_gate_exps", E * MI, H), ("ffn_up_exps", E * MI, H), ("ffn_down_exps", E * H, MI)): + raw[name] = torch.randint(0, 256, (rows, k // 32 * 34), dtype=torch.uint8) + tensors[f"blk.{blk}.{name}.weight"] = (Q8_0, (k, rows // E, E), raw[name].numpy().tobytes()) + else: + for name, shape in (("ffn_gate_exps", (E, MI, H)), ("ffn_up_exps", (E, MI, H)), ("ffn_down_exps", (E, H, MI))): + put(name, torch.randn(*shape)) + else: + for name, shape in (("ffn_gate", (I, H)), ("ffn_up", (I, H)), ("ffn_down", (H, I))): + put(name, torch.randn(*shape)) + put("nextn.eh_proj", torch.randn(H, 2 * H)) + for name in ("nextn.enorm", "nextn.hnorm", "nextn.shared_head_norm"): + put(name, torch.randn(H), F32) + kv = {"block_count": blk + nextn_layers, "embedding_length": H} + if nextn_layers: + kv["nextn_predict_layers"] = nextn_layers + path = tmp_path / f"trunk_{'moe' if moe else 'dense'}.gguf" + _write_gguf(path, "qwen35moe" if moe else "qwen35", kv, tensors) + return str(path), raw, q8 + + +def _gguf_draft(path: str, *, moe: bool, tmp_path): + """The draft the engine builds for the GGUF trunk (its quant the file's ggml types), and its state dict.""" + from dataclasses import replace + + from freetoken.layers import rotary + from freetoken.models.qwen3_5_moe.mtp import Qwen3_5MTPDraft + from freetoken.spec.mtp.proposer import with_resident_moe + from freetoken.utils.torch_utils import torch_dtype + + spec = get_model_spec("Qwen3_5MoeGGUFForCausalLM" if moe else "Qwen3_5GGUFForCausalLM") + quant = checkpoint_quant_config(path, SimpleNamespace(model_path=path), spec) + folder = tmp_path / "cfg" + folder.mkdir(exist_ok=True) + (folder / "config.json").write_text(json.dumps(_config_json(moe, None))) + trunk = replace(parse_config(cached_load_hf_config(str(folder))), quant=quant) + saved = rotary._ROPE_DEVICE + rotary.set_rope_device(torch.device("cpu")) + rotary.get_rope.cache_clear() + try: + with torch.device("meta"), torch_dtype(torch.bfloat16): + draft = Qwen3_5MTPDraft(with_resident_moe(trunk), SimpleNamespace(model=SimpleNamespace(embed_tokens=None))) + finally: + rotary.set_rope_device(saved) + rotary.get_rope.cache_clear() + return draft, draft.state_dict() + + +@pytest.mark.parametrize("moe", [False, True]) +def test_nextn_block_maps_onto_the_draft_state_dict(tmp_path, moe): + """The nextn reader fills exactly the draft the safetensors path builds: same names for an all-float file, norms un-plus-one'd, projections fused in q|k|v and gate|up order.""" + from freetoken.models.qwen3_5_moe.gguf import iter_nextn_weights, nextn_block_index + from freetoken.models.qwen3_5_moe.mtp import Qwen3_5MTPDraft + + path, raw, _ = _nextn_gguf(tmp_path, moe=moe) + Qwen3_5MTPDraft.check_source(path, path) + assert nextn_block_index(path) == 2 + draft, state = _gguf_draft(path, moe=moe, tmp_path=tmp_path) + loaded = {n: t.to(dtype=state[n].dtype) for n, t in iter_nextn_weights(path, draft.config)} + assert set(loaded) == set(state) + for key, tensor in loaded.items(): + assert tensor.shape == state[key].shape and tensor.dtype is state[key].dtype, key + + def norm(name: str) -> torch.Tensor: + return (raw[name] - 1.0).to(torch.bfloat16) + + for key, source in ( + ("norm.weight", "nextn.shared_head_norm"), ("pre_fc_norm_embedding.weight", "nextn.enorm"), + ("pre_fc_norm_hidden.weight", "nextn.hnorm"), ("layers.0.input_layernorm.weight", "attn_norm"), + ("layers.0.post_attention_layernorm.weight", "post_attention_norm"), + ("layers.0.self_attn.q_norm.weight", "attn_q_norm"), ("layers.0.self_attn.k_norm.weight", "attn_k_norm"), + ): + assert torch.equal(loaded[key], norm(source)), key + assert _same(loaded["fc.weight"], raw["nextn.eh_proj"].to(torch.bfloat16)) + qkv = loaded["layers.0.self_attn.qkv_proj.weight"] + for part, source in zip(_slices(qkv, [raw[n] for n in ("attn_q", "attn_k", "attn_v")]), ("attn_q", "attn_k", "attn_v")): + assert _same(part, raw[source].to(torch.bfloat16)) + assert _same(loaded["layers.0.self_attn.o_proj.weight"], raw["attn_output"].to(torch.bfloat16)) + if moe: + gate_up = loaded["layers.0.mlp.shared_expert.gate_up_proj.weight"] + assert _same(gate_up[:I], raw["ffn_gate_shexp"].to(torch.bfloat16)) + assert _same(gate_up[I:], raw["ffn_up_shexp"].to(torch.bfloat16)) + assert torch.equal(loaded["layers.0.mlp.shared_expert_gate.weight"], raw["ffn_gate_inp_shexp"].reshape(1, -1).to(torch.bfloat16)) + assert not any(".experts." in k for k in loaded) + else: + gate_up = loaded["layers.0.mlp.gate_up_proj.weight"] + assert _same(gate_up[:I], raw["ffn_gate"].to(torch.bfloat16)) + assert _same(gate_up[I:], raw["ffn_up"].to(torch.bfloat16)) + assert _same(loaded["layers.0.mlp.down_proj.weight"], raw["ffn_down"].to(torch.bfloat16)) + + +def test_a_quantized_nextn_tensor_lands_in_its_own_state_dict_leaf(tmp_path): + """A qkv whose q half is ggml-quantized and k|v float fills one packed ``weight`` leaf and one float ``weight_1`` leaf, bytes as stored.""" + from freetoken.layers.quantization.ggml_types import pad_rows, row_bytes + from freetoken.models.qwen3_5_moe.gguf import iter_nextn_weights + + path, raw, q8 = _nextn_gguf(tmp_path, moe=False, q_quantized=True) + draft, state = _gguf_draft(path, moe=False, tmp_path=tmp_path) + loaded = {n: t.to(dtype=state[n].dtype) for n, t in iter_nextn_weights(path, draft.config)} + for key, tensor in loaded.items(): + assert tensor.shape == state[key].shape and tensor.dtype is state[key].dtype, key + packed, float_rows = loaded["layers.0.self_attn.qkv_proj.weight"], loaded["layers.0.self_attn.qkv_proj.weight_1"] + rows = 2 * QH * AHD + assert packed.dtype is torch.uint8 and packed.shape == (rows + pad_rows(Q8_0, H), row_bytes(Q8_0, H)) + assert _same(packed[:rows], q8) + assert float_rows.dtype is torch.bfloat16 + for part, source in zip(_slices(float_rows, [raw[n] for n in ("attn_k", "attn_v")]), ("attn_k", "attn_v")): + assert _same(part, raw[source].to(torch.bfloat16)) + # the eh_proj of the same file stays float, one plain leaf + assert _same(loaded["fc.weight"], raw["nextn.eh_proj"].to(torch.bfloat16)) + + +def test_nextn_expert_banks_fill_the_drafts_resident_moe(tmp_path): + """The nextn block's stacked experts become the draft's one resident ggml bank set, gate rows before up rows.""" + from freetoken.layers.moe import OffloadMoELayer, iter_moe_layers + from freetoken.models.qwen3_5_moe.gguf import nextn_expert_banks + from freetoken.moe.expert_banks import attach_resident_banks + + path, raw, _ = _nextn_gguf(tmp_path, moe=True) + draft, _state = _gguf_draft(path, moe=True, tmp_path=tmp_path) + moe = draft.layers.op_list[0].mlp.experts + assert not any(isinstance(m, OffloadMoELayer) for m in iter_moe_layers(draft)) + banks = nextn_expert_banks(path, torch.device("cpu")) + attach_resident_banks([moe], banks) + gate_up, down = moe._expert_banks["gate_up"], moe._expert_banks["down"] + assert gate_up.dtype is torch.bfloat16 and gate_up.shape == (E, 2 * MI, H) + assert _same(gate_up[:, :MI], raw["ffn_gate_exps"].to(torch.bfloat16)) + assert _same(gate_up[:, MI:], raw["ffn_up_exps"].to(torch.bfloat16)) + assert _same(down, raw["ffn_down_exps"].to(torch.bfloat16)) + + +def test_a_no_mtp_gguf_refusal_names_the_two_causes(tmp_path): + from freetoken.models.qwen3_5_moe.mtp import Qwen3_5MTPDraft + + path, _raw, _q8 = _nextn_gguf(tmp_path, moe=True, nextn_layers=0) + with pytest.raises(ValueError, match=r"no MTP \(nextn\) block.*--no-mtp.*Strata"): + Qwen3_5MTPDraft.check_source(path, path) + + +def test_the_sidecar_draft_loads_under_the_sidecars_own_quant(tmp_path): + """`load` builds the draft's quant from the sidecar's own tensor table: the target's class + is the one the HF and GGUF rows share, and that name's HF row builds the draft unquantized + against the sidecar's packed rows.""" + from freetoken.layers import rotary + from freetoken.layers.quantization.ggml_types import pad_rows, row_bytes + from freetoken.models.qwen3_5_moe.mtp import Qwen3_5MTPDraft + from freetoken.spec.mtp.proposer import with_resident_moe + + for name in ("t", "q8"): + (tmp_path / name).mkdir() + trunk, _raw, _ = _nextn_gguf(tmp_path / "t", moe=False) + sidecar, _r, q8 = _nextn_gguf(tmp_path / "q8", moe=False, q_quantized=True) + (tmp_path / "cfg").mkdir() + (tmp_path / "cfg" / "config.json").write_text(json.dumps(_config_json(False, None))) + trunk_config = parse_config(cached_load_hf_config(str(tmp_path / "cfg"))) + # the fake target carries the class name the HF and GGUF rows share + target = type("Qwen3_5ForConditionalGeneration", (), {})() + target.model = SimpleNamespace(embed_tokens=None) + + saved = rotary._ROPE_DEVICE + rotary.set_rope_device(torch.device("cpu")) + rotary.get_rope.cache_clear() + try: + draft = Qwen3_5MTPDraft.load(with_resident_moe(trunk_config), target, sidecar, torch.device("cpu"), torch.bfloat16) + finally: + rotary.set_rope_device(saved) + rotary.get_rope.cache_clear() + + state = draft.state_dict() + rows = 2 * QH * AHD + packed, float_rows = state["layers.0.self_attn.qkv_proj.weight"], state["layers.0.self_attn.qkv_proj.weight_1"] + assert packed.dtype is torch.uint8 and packed.shape == (rows + pad_rows(Q8_0, H), row_bytes(Q8_0, H)) + assert _same(packed[:rows], q8) + assert float_rows.dtype is torch.bfloat16 and float_rows.shape == (2 * KVH * AHD, H) + # eh_proj stays float in the same file: one plain bf16 leaf + assert state["fc.weight"].dtype is torch.bfloat16 + + +def test_quantized_nextn_expert_banks_keep_their_packed_rows(tmp_path): + """The nextn experts at their shipped ggml types stay packed: gate rows before up rows + per expert, byte for byte as stored, the down bank in its own row width.""" + from freetoken.layers.moe import iter_moe_layers, OffloadMoELayer + from freetoken.layers.quantization.ggml_types import row_bytes + from freetoken.models.qwen3_5_moe.gguf import nextn_expert_banks + from freetoken.moe.expert_banks import attach_resident_banks + + path, raw, _ = _nextn_gguf(tmp_path, moe=True, experts_quantized=True) + draft, _state = _gguf_draft(path, moe=True, tmp_path=tmp_path) + moe = draft.layers.op_list[0].mlp.experts + assert not any(isinstance(m, OffloadMoELayer) for m in iter_moe_layers(draft)) + banks = nextn_expert_banks(path, torch.device("cpu")) + attach_resident_banks([moe], banks) + gate_up, down = moe._expert_banks["gate_up"], moe._expert_banks["down"] + rb, rb_down = row_bytes(Q8_0, H), row_bytes(Q8_0, MI) + assert gate_up.dtype is torch.uint8 and gate_up.shape == (E, 2 * MI, rb) + assert down.dtype is torch.uint8 and down.shape == (E, H, rb_down) + for e in range(E): + assert torch.equal(gate_up[e, :MI].reshape(-1, rb), raw["ffn_gate_exps"][e * MI : (e + 1) * MI]) + assert torch.equal(gate_up[e, MI:].reshape(-1, rb), raw["ffn_up_exps"][e * MI : (e + 1) * MI]) + assert torch.equal(down[e].reshape(-1, rb_down), raw["ffn_down_exps"][e * H : (e + 1) * H])