From 953659dbbba6cf2f822a25787027c6210718a8d0 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Thu, 10 Sep 2026 22:36:28 +0000 Subject: [PATCH 1/2] feat(qwen3_5_moe): read every checkpoint layout through the QuantConfig NVFP4 dense layers now declare the input_scale their scheme carries, so an FTW written by an older build from an NVFP4 dense checkpoint needs scripts/ftw_hotfix.py before it loads again. --- .../layers/quantization/configs/base.py | 20 + .../configs/compressed_tensors.py | 13 +- .../layers/quantization/configs/fp8.py | 16 +- .../layers/quantization/configs/modelopt.py | 10 +- .../layers/quantization/configs/mxfp4.py | 6 +- .../layers/quantization/linear/nvfp4.py | 3 +- python/freetoken/layers/quantization/names.py | 12 +- python/freetoken/models/gemma4/weight.py | 35 +- python/freetoken/models/glm4_moe/weight.py | 2 + python/freetoken/models/glm5_next/config.py | 30 +- python/freetoken/models/loader.py | 65 +- .../freetoken/models/muse_glimmer/weight.py | 4 +- python/freetoken/models/nvfp4_banks.py | 5 +- python/freetoken/models/qwen3_5_moe/config.py | 145 +-- python/freetoken/models/qwen3_5_moe/weight.py | 969 +++++------------- python/freetoken/models/register.py | 4 + python/freetoken/models/weight.py | 15 +- tests/models/test_muse_glimmer.py | 4 +- tests/models/test_quant_config.py | 67 ++ tests/models/test_qwen3_5_moe_weight.py | 553 ++++++++++ 20 files changed, 1030 insertions(+), 948 deletions(-) create mode 100644 tests/models/test_qwen3_5_moe_weight.py diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py index d786a3338..2a4ef432e 100644 --- a/python/freetoken/layers/quantization/configs/base.py +++ b/python/freetoken/layers/quantization/configs/base.py @@ -3,6 +3,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import Any, ClassVar from ..linear import LinearConfig @@ -36,8 +37,18 @@ def quantization_config_of(hf_config: Any) -> dict[str, Any] | None: return dict(vars(q)) +@dataclass(frozen=True) +class Stored: + """One checkpoint tensor behind a role: its suffix, and whether it holds the quant-side scale whose reciprocal the layer wants.""" + + name: str + reciprocal: bool = False + + class QuantConfig(ABC): dialect: ClassVar[str] + # per kind the dialect exports, role -> the checkpoint tensor suffix (or Stored) that feeds it + STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] def __init__(self, name_map: NameMap | None = None, unquantized: tuple[str, ...] = ()): self.name_map = name_map or NameMap() @@ -65,6 +76,14 @@ def scheme_for(self, prefix: str) -> QuantScheme | None: self._schemes[prefix] = scheme return scheme + def stored_tensors(self, kind: QuantKind) -> dict[str, Stored]: + """role -> checkpoint tensor for every role the dialect stores for ``kind``.""" + return {role: entry if isinstance(entry, Stored) else Stored(entry) for role, entry in self.STORAGE[kind].items()} + + def storage(self, scheme: QuantScheme) -> dict[str, Stored]: + """role -> checkpoint tensor for one scheme's tensors.""" + return {role: entry for role, entry in self.stored_tensors(scheme.kind).items() if scheme.has(role)} + def get_quant_method(self, layer: Any, prefix: str): scheme = self.scheme_for(prefix) layer_kind = layer.quant_layer_kind @@ -101,6 +120,7 @@ class NoQuantConfig(QuantConfig): """No ``quantization_config``: every module is bf16.""" dialect = "none" + STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = {} @classmethod def claims(cls, q: dict[str, Any]) -> bool: diff --git a/python/freetoken/layers/quantization/configs/compressed_tensors.py b/python/freetoken/layers/quantization/configs/compressed_tensors.py index d2b31dbac..a0af4beeb 100644 --- a/python/freetoken/layers/quantization/configs/compressed_tensors.py +++ b/python/freetoken/layers/quantization/configs/compressed_tensors.py @@ -4,9 +4,9 @@ from ..names import Matcher, ct_set from ..registry import register_dialect -from ..scheme import QuantScheme +from ..scheme import QuantKind, QuantScheme from ..scheme import fp8_block_scheme, fp8_tensor_scheme, nvfp4_scheme -from .base import QuantConfig +from .base import QuantConfig, Stored @register_dialect @@ -23,6 +23,15 @@ class CompressedTensorsConfig(QuantConfig): "FP8_CHANNEL": fp8_tensor_scheme("bf16", per_row=True), "FP8_BLOCK": fp8_block_scheme("float"), } + # the NVFP4 globals are the quant-side scales (vLLM: alpha = 1 / (input_global_scale * weight_global_scale)) + STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = { + QuantKind.NVFP4: { + "weight": "weight_packed", "weight_scale": "weight_scale", + "weight_global": Stored("weight_global_scale", reciprocal=True), "input_scale": Stored("input_global_scale", reciprocal=True), + }, + QuantKind.FP8_TENSOR: {"weight": "weight", "weight_scale": "weight_scale", "input_scale": "input_scale"}, + QuantKind.FP8_BLOCK: {"weight": "weight", "weight_scale_inv": "weight_scale"}, + } def __init__(self, q: dict[str, Any], hf_config: Any = None, *, name_map=None, unquantized=()): super().__init__(name_map, unquantized) diff --git a/python/freetoken/layers/quantization/configs/fp8.py b/python/freetoken/layers/quantization/configs/fp8.py index 11c8fbe20..c23cbfdee 100644 --- a/python/freetoken/layers/quantization/configs/fp8.py +++ b/python/freetoken/layers/quantization/configs/fp8.py @@ -4,9 +4,9 @@ from ..names import is_routed_expert, name_set, substr_set from ..registry import register_dialect -from ..scheme import QuantScheme +from ..scheme import QuantKind, QuantScheme from ..scheme import FP8_BLOCK, fp8_block_scheme, fp8_tensor_scheme, mxfp4_scheme -from .base import QuantConfig, cfg_get +from .base import QuantConfig, Stored, cfg_get @register_dialect @@ -22,6 +22,12 @@ class Fp8BlockConfig(QuantConfig): "TABLE": fp8_tensor_scheme("float"), "EXPERT_MXFP4": mxfp4_scheme(), } + # transformers' fp8 names; DeepSeek-V4's e8m0 export calls every scale ``scale`` (see storage) + STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = { + QuantKind.FP8_BLOCK: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"}, + QuantKind.FP8_TENSOR: {"weight": "weight", "weight_scale": "weight_scale"}, + QuantKind.MXFP4: {"weight": "weight", "weight_scale": "scale"}, + } def __init__(self, q: dict[str, Any], hf_config: Any = None, *, name_map=None, unquantized=()): super().__init__(name_map, unquantized) @@ -36,6 +42,12 @@ def __init__(self, q: dict[str, Any], hf_config: Any = None, *, name_map=None, u self.e8m0 = str(q.get("scale_fmt") or "").lower() == "ue8m0" self.expert_fp4 = str(cfg_get(hf_config, "expert_dtype") or "").lower() == "fp4" + def storage(self, scheme: QuantScheme) -> dict[str, Stored]: + names = super().storage(scheme) + if self.e8m0 and scheme.kind is QuantKind.FP8_BLOCK: + names["weight_scale_inv"] = Stored("scale") + return names + def scheme_for_name(self, name: str) -> QuantScheme | None: if self.convert_tables(name): return self.SCHEMES["TABLE"] diff --git a/python/freetoken/layers/quantization/configs/modelopt.py b/python/freetoken/layers/quantization/configs/modelopt.py index 813b0667f..1d89b2358 100644 --- a/python/freetoken/layers/quantization/configs/modelopt.py +++ b/python/freetoken/layers/quantization/configs/modelopt.py @@ -4,9 +4,9 @@ from ..names import ancestors, name_set from ..registry import register_dialect -from ..scheme import QuantScheme +from ..scheme import QuantKind, QuantScheme from ..scheme import fp8_block_scheme, fp8_tensor_scheme, mxfp8_scheme, nvfp4_scheme -from .base import QuantConfig +from .base import QuantConfig, Stored @register_dialect @@ -24,6 +24,12 @@ class ModelOptConfig(QuantConfig): "FP8_PB_WO": fp8_block_scheme("fp32"), "MXFP8": mxfp8_scheme(), } + STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = { + QuantKind.FP8_TENSOR: {"weight": "weight", "weight_scale": "weight_scale", "input_scale": "input_scale"}, + QuantKind.FP8_BLOCK: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"}, + QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"}, + QuantKind.NVFP4: {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2", "input_scale": "input_scale"}, + } @classmethod def claims(cls, q: dict[str, Any]) -> bool: diff --git a/python/freetoken/layers/quantization/configs/mxfp4.py b/python/freetoken/layers/quantization/configs/mxfp4.py index e9458f034..5eb24388d 100644 --- a/python/freetoken/layers/quantization/configs/mxfp4.py +++ b/python/freetoken/layers/quantization/configs/mxfp4.py @@ -4,9 +4,9 @@ from ..names import name_set from ..registry import register_dialect -from ..scheme import QuantScheme +from ..scheme import QuantKind, QuantScheme from ..scheme import mxfp4_scheme -from .base import QuantConfig +from .base import QuantConfig, Stored @register_dialect @@ -16,6 +16,8 @@ class Mxfp4Config(QuantConfig): dialect = "mxfp4" SCHEME: ClassVar[QuantScheme] = mxfp4_scheme() + # the experts are stacked per layer (``gate_up_proj_blocks`` / ``_scales``), not per-Linear tensors; gpt_oss reads them itself + STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = {} def __init__(self, q: dict[str, Any], hf_config: Any = None, *, name_map=None, unquantized=()): super().__init__(name_map, unquantized) diff --git a/python/freetoken/layers/quantization/linear/nvfp4.py b/python/freetoken/layers/quantization/linear/nvfp4.py index 6644acdb2..0953b8250 100644 --- a/python/freetoken/layers/quantization/linear/nvfp4.py +++ b/python/freetoken/layers/quantization/linear/nvfp4.py @@ -120,4 +120,5 @@ def create_weights(self, layer: Any) -> None: layer.weight = torch.empty(g.out_features, g.in_features // 2, dtype=torch.uint8) layer.weight_scale = torch.empty(g.out_features, g.in_features // GROUP, dtype=FP8) layer.weight_global = torch.empty(g.out_features, dtype=torch.float16) - # input_scale stays undeclared: the W4A16 kernels never read it and today's readers drop it + # no W4A16 kernel reads it; declared so a W4A4 checkpoint loads complete and a W4A4 kernel finds it in place + layer.input_scale = torch.empty((), dtype=torch.float32) if self.scheme.has("input_scale") else None diff --git a/python/freetoken/layers/quantization/names.py b/python/freetoken/layers/quantization/names.py index c7a79743c..2cc47fabf 100644 --- a/python/freetoken/layers/quantization/names.py +++ b/python/freetoken/layers/quantization/names.py @@ -38,12 +38,18 @@ def substr_set(patterns: tuple[str, ...]) -> Matcher: def ct_set(patterns: tuple[str, ...], *, class_names: bool) -> Matcher: - """compressed-tensors ``targets`` / ``ignore``: module names, ``re:`` regexes, or the class name Linear.""" - names = name_set(tuple(p for p in patterns if not p.startswith("re:") and p != "Linear")) + """compressed-tensors ``targets`` / ``ignore``: module names, ``re:`` regexes, or the class name Linear. + + A name covers that module alone, not its children: llm-compressor lists every skipped module, containers included, so an ``ignore`` entry for ``linear_attn`` says nothing about ``linear_attn.in_proj_qkv``.""" + names = frozenset(p for p in patterns if not p.startswith("re:") and p != "Linear") + # a class name other than Linear cannot be matched from a module name alone; fail here rather than serve the module bf16 + unknown = [p for p in names if "." not in p and p[:1].isupper()] if class_names else [] + if unknown: + raise NotImplementedError(f"compressed-tensors target class {unknown[0]!r} is not supported; only Linear is") regexes = [p[3:] for p in patterns if p.startswith("re:")] rx = re.compile("|".join(f"(?:{r})" for r in regexes)) if regexes else None any_linear = class_names and "Linear" in patterns - return lambda name: any_linear or names(name) or (rx is not None and rx.match(name) is not None) + return lambda name: any_linear or name in names or (rx is not None and rx.match(name) is not None) _ROUTED_EXPERT = re.compile(r"\.experts\.\d+(\.|$)") diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index 71f5531b2..0d2056a63 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -59,16 +59,13 @@ # modelopt-NVFP4 dense MLP (nvidia/Gemma-4-31B-IT-NVFP4): mlp.{gate,up,down}_proj are W4A16 # FP4 -- uint8 weight + fp8-e4m3 block weight_scale + per-tensor weight_scale_2 + input_scale. -# The scales are consumed with their .weight; input_scale is unused (W4A16). Mirrors the -# qwen3_5_moe native-NVFP4 dense loader. +# The scales are consumed with their .weight. _NVFP4_DENSE_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") _NVFP4_DENSE_MLP_RE = re.compile(r"\.mlp\.(gate_proj|up_proj|down_proj)\.weight$") -def _nvfp4_dense_parts(f, raw_base: str): - """Load an NVFP4 dense weight as the W4A16 kernel's buffers: (weight uint8 [O, IN//2], - weight_scale fp8-e4m3 block [O, IN//16], weight_global fp16 [O] from the per-tensor - weight_scale_2 broadcast per output row).""" +def _nvfp4_dense_parts(f, raw_base: str, keyset: set[str]): + """Load an NVFP4 dense weight as the NVFP4 linear method's buffers: weight uint8 [O, IN//2], weight_scale fp8-e4m3 block [O, IN//16], weight_global fp16 [O] (the per-tensor weight_scale_2 per output row), input_scale fp32 scalar or None when the export has none.""" w = f.get_tensor(raw_base + ".weight") s = f.get_tensor(raw_base + ".weight_scale") g = f.get_tensor(raw_base + ".weight_scale_2").reshape(1).to(torch.float16) @@ -78,31 +75,37 @@ def _nvfp4_dense_parts(f, raw_base: str): and s.dtype is torch.float8_e4m3fn and g.dtype is torch.float16 ), f"unexpected NVFP4 dense dtypes at {raw_base}: {w.dtype}/{s.dtype}/{g.dtype}" - return w, s, g + a = f.get_tensor(raw_base + ".input_scale").reshape(()).to(torch.float32) if raw_base + ".input_scale" in keyset else None + return w, s, g, a -def _emit_nvfp4_dense_mlp(f, base: str, raw_base: str, buf: dict): - """(key, tensor) triples for an NVFP4 dense MLP projection: down_proj standalone; +def _emit_nvfp4_dense_mlp(f, base: str, raw_base: str, buf: dict, keyset: set[str]): + """(key, tensor) pairs for an NVFP4 dense MLP projection: down_proj standalone; gate_proj/up_proj merged output-wise into gate_up_proj (each keeps its own scales, so the fused weight is exact). Returns [] while a gate/up merge is still buffered.""" - w, s, g = _nvfp4_dense_parts(f, raw_base) + w, s, g, a = _nvfp4_dense_parts(f, raw_base, keyset) if base.endswith(".down_proj"): - return [(base + ".weight", w), (base + ".weight_scale", s), (base + ".weight_global", g)] + out = [(base + ".weight", w), (base + ".weight_scale", s), (base + ".weight_global", g)] + return out + ([(base + ".input_scale", a)] if a is not None else []) is_gate = base.endswith(".gate_proj") prefix = base[: -len(".gate_proj")] if is_gate else base[: -len(".up_proj")] slots = buf.setdefault(prefix, {}) - slots["gate" if is_gate else "up"] = (w, s, g) + slots["gate" if is_gate else "up"] = (w, s, g, a) if "gate" not in slots or "up" not in slots: return [] - gw, gs, gg = slots["gate"] - uw, us, ug = slots["up"] + gw, gs, gg, ga = slots["gate"] + uw, us, ug, ua = slots["up"] del buf[prefix] pre = prefix + ".gate_up_proj" - return [ + out = [ (pre + ".weight", torch.cat([gw, uw], dim=0)), (pre + ".weight_scale", torch.cat([gs, us], dim=0)), (pre + ".weight_global", torch.cat([gg, ug], dim=0)), ] + if ga is not None and ua is not None: + # both parts read the same activation; the larger range covers both + out.append((pre + ".input_scale", torch.maximum(ga, ua))) + return out def _rename_language_key(raw_name: str) -> str: @@ -207,7 +210,7 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: and raw_name[: -len(".weight")] + ".weight_scale_2" in keyset ): yield from _emit_nvfp4_dense_mlp( - f, name[: -len(".weight")], raw_name[: -len(".weight")], gateup_buf + f, name[: -len(".weight")], raw_name[: -len(".weight")], gateup_buf, keyset ) continue diff --git a/python/freetoken/models/glm4_moe/weight.py b/python/freetoken/models/glm4_moe/weight.py index 5d4f9fcea..38b17a053 100644 --- a/python/freetoken/models/glm4_moe/weight.py +++ b/python/freetoken/models/glm4_moe/weight.py @@ -87,6 +87,8 @@ def _iter_nvfp4_resident( yield f"{dst_prefix}.weight", packed yield f"{dst_prefix}.weight_scale", scale yield f"{dst_prefix}.weight_global", g.expand(packed.shape[0]).contiguous() + if reader.has(f"{src_prefix}.input_scale"): + yield f"{dst_prefix}.input_scale", reader.get(f"{src_prefix}.input_scale").reshape(()).to(torch.float32) def _iter_attn_df11( diff --git a/python/freetoken/models/glm5_next/config.py b/python/freetoken/models/glm5_next/config.py index 1cc6d798e..349584667 100644 --- a/python/freetoken/models/glm5_next/config.py +++ b/python/freetoken/models/glm5_next/config.py @@ -48,6 +48,34 @@ def _dsa_on(args, dsa_layer_ids) -> bool: ) +def _quant_accessor(hf_config: Any): + """A ``get(key, default=None)`` accessor over the HF ``quantization_config`` (dict or + object), or ``None`` when the model has no quant config.""" + quant = getattr(hf_config, "quantization_config", None) + if quant is None: + return None + return quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d)) + + +def _fp8_block_quant(hf_config: Any) -> tuple[str, tuple[int, int] | None]: + """Detect DeepSeek-V3-style 128x128 block-fp8 from HF ``quantization_config``. + + Returns ``("fp8_block", (block_n, block_k))`` for a block-fp8 checkpoint (weights + fp8-e4m3 + per-block ``weight_scale_inv``, dynamic activation), else ``("none", None)``. + The quantization_config sits on the top-level hf_config (not ``text_config``). + """ + get = _quant_accessor(hf_config) + if get is None: + return "none", None + method = str(get("quant_method") or get("quant_algo") or "").lower() + block = get("weight_block_size") + if method == "fp8" and block: + bs = tuple(int(x) for x in block) + assert bs == (128, 128), f"only 128x128 block-fp8 is supported, got {bs}" + return "fp8_block", bs + return "none", None + + def parse_config(hf_config: Any) -> ModelConfig: args = load_args(hf_config) text = getattr(hf_config, "text_config", hf_config) @@ -128,8 +156,6 @@ def parse_config(hf_config: Any) -> ModelConfig: t == "sparse" for t in mlp_types[first_dense:] ), f"mlp_layer_types is not a dense-prefix layout: {mlp_types}" - from freetoken.models.qwen3_5_moe.config import _fp8_block_quant # shared until the legacy quant fields go - expert_quant, weight_block_size = _fp8_block_quant(hf_config) if expert_quant == "none": expert_quant = detect_expert_quant(hf_config) diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 0a2b69aee..0f310ccbb 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -4,6 +4,7 @@ import json import os import re +import struct from dataclasses import dataclass from typing import Iterable, Iterator @@ -53,6 +54,23 @@ def iter_weight_files(model_path: str) -> list[str]: return [f for f in files if not f.endswith("consolidated.safetensors")] or files +def safetensors_weight_map(folder: str) -> dict[str, str]: + """Tensor name -> shard basename, from the index or from each shard's header when the checkpoint ships none.""" + index = os.path.join(folder, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as f: + return json.load(f)["weight_map"] + weight_map: dict[str, str] = {} + for path in sorted(iter_weight_files(folder)): + with open(path, "rb") as fh: + n = struct.unpack(" None: """drop a file's page cache: banks + full checkpoint cache don't both fit in host RAM (OOM).""" try: @@ -135,14 +153,13 @@ def iter_merged_tensors( # --------------------------------------------------------------------------------- # compressed-tensors NVFP4 (llm-compressor) dense-weight helpers, shared by the -# models that serve such checkpoints natively (qwen3_5_moe, muse_glimmer). Storage: +# models that serve such checkpoints natively (muse_glimmer). Storage: # ``weight_packed`` (uint8 [O, IN//2]) + ``weight_scale`` (fp8-e4m3 block [O, IN//16]) # + a scalar ``weight_global_scale``. The stored global is the *quant-side* scale, so # the dequant/native global is its reciprocal (vLLM inverts it identically). # --------------------------------------------------------------------------------- -# Quant scales consumed with their ``weight_packed`` (or unused: the input scales are -# for W4A4 activation quant, which FreeToken does not run). +# Quant scales consumed with their ``weight_packed``. CT_SCALE_SUFFIXES = ( ".weight_scale", ".weight_global_scale", ".input_global_scale", ".input_scale", ) @@ -158,27 +175,16 @@ class ShardReader: def __init__(self, model_path: str, device: torch.device): folder = download_hf_weight(model_path) - index = os.path.join(folder, "model.safetensors.index.json") - if os.path.exists(index): - with open(index, encoding="utf-8") as f: - weight_map = json.load(f)["weight_map"] - self._map = { - name: os.path.join(folder, shard) for name, shard in weight_map.items() - } - else: # single-file checkpoint - import safetensors - - self._map = {} - for file in iter_weight_files(model_path): - with safetensors.safe_open(file, framework="pt", device="cpu") as f: - for name in f.keys(): - self._map[name] = file + self._map = {name: os.path.join(folder, shard) for name, shard in safetensors_weight_map(folder).items()} self._device = str(device) self._handles: dict[str, object] = {} def files(self) -> list[str]: return sorted(set(self._map.values())) + def has(self, name: str) -> bool: + return name in self._map + def names_in(self, file: str) -> list[str]: return [name for name, shard in self._map.items() if shard == file] @@ -201,21 +207,21 @@ def close(self) -> None: self._handles.clear() -def nvfp4_parts_ct(f, raw_base: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """compressed-tensors NVFP4 -> ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], - per-output-row global fp16 [O])`` for the W4A16 kernels.""" +def nvfp4_parts_ct(f, raw_base: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """compressed-tensors NVFP4 -> ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], per-output-row global fp16 [O], dequant-side input_scale fp32 scalar or None)`` for the NVFP4 linear method.""" w = f.get_tensor(raw_base + ".weight_packed") s = f.get_tensor(raw_base + ".weight_scale") wg = f.get_tensor(raw_base + ".weight_global_scale").reshape(1).to(torch.float32) g = (1.0 / wg).to(torch.float16).expand(w.shape[0]).contiguous() - return w, s, g + a = None + if f.has(raw_base + ".input_global_scale"): + a = (1.0 / f.get_tensor(raw_base + ".input_global_scale").reshape(()).to(torch.float32)) + return w, s, g, a def ct_nvfp4_fuse(base: str, parts_tuple: tuple, buf: dict, groups: dict[str, tuple[str, ...]]): - """Buffer a native NVFP4 fusion part ``(w, s, g)``; emit the concatenated native parts - (``.weight``/``.weight_scale``/``.weight_global``, output-dim concat with each part - keeping its own scales, so the fused FP4 weight is exact) once complete, ``[]`` while - incomplete, ``None`` if ``base`` is not a fusion part of any group in ``groups``.""" + """Buffer a native NVFP4 fusion part ``(w, s, g, a)``; once complete, emit the concatenated ``.weight`` / ``.weight_scale`` / ``.weight_global`` (output-dim concat, each part keeps its own scales, so the fused FP4 weight is exact) plus the largest ``.input_scale`` when every part has one. + ``[]`` while incomplete, ``None`` if ``base`` is not a fusion part of any group in ``groups``.""" for fused_suffix, parts in groups.items(): for idx, part in enumerate(parts): if base.endswith(part): @@ -228,11 +234,15 @@ def ct_nvfp4_fuse(base: str, parts_tuple: tuple, buf: dict, groups: dict[str, tu ws = [slots[i][0] for i in range(len(parts))] ss = [slots[i][1] for i in range(len(parts))] gs = [slots[i][2] for i in range(len(parts))] - return [ + acts = [slots[i][3] for i in range(len(parts))] + out = [ (key + ".weight", torch.cat(ws, dim=0)), (key + ".weight_scale", torch.cat(ss, dim=0)), (key + ".weight_global", torch.cat(gs, dim=0)), ] + if all(a is not None for a in acts): + out.append((key + ".input_scale", torch.stack(acts).max())) + return out return None @@ -293,6 +303,7 @@ def iter_stacked_experts( "MergeRule", "iter_root_safetensor_files_from_index", "iter_weight_files", + "safetensors_weight_map", "iter_merged_tensors", "iter_stacked_experts", "shard_tensor", diff --git a/python/freetoken/models/muse_glimmer/weight.py b/python/freetoken/models/muse_glimmer/weight.py index 8a6132254..bf82b3ab8 100644 --- a/python/freetoken/models/muse_glimmer/weight.py +++ b/python/freetoken/models/muse_glimmer/weight.py @@ -123,10 +123,12 @@ def _iter_weights_compressed_tensors( if emit is not None: yield from emit else: # standalone: o_proj, down_proj - w, s, g = parts + w, s, g, a = parts yield base + ".weight", w yield base + ".weight_scale", s yield base + ".weight_global", g + if a is not None: + yield base + ".input_scale", a continue yield name, reader.get_tensor(raw_name) finally: diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 1d1b42859..2dfdaaa35 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -1,7 +1,6 @@ from __future__ import annotations import collections -import json import os import re from dataclasses import dataclass @@ -83,12 +82,12 @@ def iter_nvfp4_expert_pieces( way tensors of one expert may span shards, so they are grouped by (layer, expert) as they land. """ from freetoken.models.loader import drop_page_cache as _drop + from freetoken.models.loader import safetensors_weight_map from freetoken.moe.expert_pieces import per_expert_pieces drop = drop_page_cache or _drop folder = download_hf_weight(model_path) - with open(os.path.join(folder, "model.safetensors.index.json"), encoding="utf-8") as f: - weight_map = json.load(f)["weight_map"] + weight_map = safetensors_weight_map(folder) wanted: dict[str, tuple[int, int, str]] = {} for name in weight_map: diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index dbcabc74b..cc5a62b91 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/python/freetoken/models/qwen3_5_moe/config.py @@ -2,125 +2,24 @@ from typing import Any +from freetoken.layers.quantization import QuantConfig, QuantKind from freetoken.models.config import ( FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, ModelConfig, RotaryConfig, - detect_compressed_tensors_nvfp4, ) -def _quant_accessor(hf_config: Any): - """A ``get(key, default=None)`` accessor over the HF ``quantization_config`` (dict or - object), or ``None`` when the model has no quant config.""" - quant = getattr(hf_config, "quantization_config", None) - if quant is None: - return None - return quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d)) - - -def _fp8_block_quant(hf_config: Any) -> tuple[str, tuple[int, int] | None]: - """Detect DeepSeek-V3-style 128x128 block-fp8 from HF ``quantization_config``. - - Returns ``("fp8_block", (block_n, block_k))`` for a block-fp8 checkpoint (weights - fp8-e4m3 + per-block ``weight_scale_inv``, dynamic activation), else ``("none", None)``. - The quantization_config sits on the top-level hf_config (not ``text_config``). - """ - get = _quant_accessor(hf_config) - if get is None: +def _expert_quant(hf_config: Any, text: Any) -> tuple[str, tuple[int, int] | None]: + """The routed experts' quant kind as the engine's format tag, with the scale block of block-fp8.""" + if not (getattr(text, "num_experts", 0) or 0): return "none", None - method = str(get("quant_method") or get("quant_algo") or "").lower() - block = get("weight_block_size") - if method == "fp8" and block: - bs = tuple(int(x) for x in block) - assert bs == (128, 128), f"only 128x128 block-fp8 is supported, got {bs}" - return "fp8_block", bs - return "none", None - - -def _expert_quant(hf_config: Any) -> str: - """Quantization format of the *routed* experts (the only weights served from the - offload cache). The nvidia/modelopt checkpoints are either plain NVFP4 (``quant_algo`` - ``NVFP4``) or ``MIXED_PRECISION`` (per-layer ``quantized_layers`` map); in the mixed - case the routed experts carry their own ``W4A16_NVFP4``/``FP8`` algo. Dense quantized - weights (attention / shared expert / lm_head) are served through their own schemes.""" - get = _quant_accessor(hf_config) - if get is None: - return "none" - algo = str(get("quant_algo") or get("quant_method") or "").lower() - if "fp4" in algo: - return "nvfp4" - if "mixed" in algo: - layers = get("quantized_layers") or {} - for name, spec in (layers.items() if isinstance(layers, dict) else []): - if name.endswith(".mlp.experts") or ".mlp.experts." in name: - expert_algo = str((spec or {}).get("quant_algo", "")).lower() - if "fp4" in expert_algo: - return "nvfp4" - if "fp8" in expert_algo: - return "fp8" - return "none" - - -# Detection now lives in models/config.py (shared with muse_glimmer); weight.py imports -# it under this name. -_compressed_tensors_nvfp4 = detect_compressed_tensors_nvfp4 - - -def _lm_head_quant(hf_config: Any) -> str: - """Whether the checkpoint stores ``lm_head`` as NVFP4. modelopt MIXED_PRECISION lists it in - the per-layer ``quantized_layers`` map (``W4A16_NVFP4``); pure-NVFP4 checkpoints have no - per-layer map and leave lm_head bf16. Returns ``"nvfp4"`` or ``"none"``.""" - get = _quant_accessor(hf_config) - if get is None: - return "none" - layers = get("quantized_layers") or {} - if not isinstance(layers, dict): - return "none" - for name, spec in layers.items(): - if name == "lm_head" or name.endswith(".lm_head"): - if "fp4" in str((spec or {}).get("quant_algo", "")).lower(): - return "nvfp4" - return "none" - - -def _dense_mlp_quant(hf_config: Any) -> str: - """NVFP4 on the *dense* (non-MoE) decoder MLP. modelopt MIXED_PRECISION dense checkpoints - (e.g. Qwen3.6-27B-NVFP4) list ``.mlp.{gate,up,down}_proj`` as ``W4A16_NVFP4`` in - ``quantized_layers``; MoE checkpoints have ``.mlp.experts.*`` / ``.mlp.shared_expert.*`` - instead (covered by ``expert_quant``). ``endswith(".mlp.gate_proj")`` matches only the bare - dense MLP -- not ``.mlp.shared_expert.gate_proj`` nor ``.mlp.experts.N.gate_proj``.""" - get = _quant_accessor(hf_config) - if get is None: - return "none" - layers = get("quantized_layers") or {} - if not isinstance(layers, dict): - return "none" - for name, spec in layers.items(): - if name.endswith((".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj")): - if "fp4" in str((spec or {}).get("quant_algo", "")).lower(): - return "nvfp4" - return "none" - - -def _attn_quant(hf_config: Any) -> str: - """Per-tensor FP8 on the *dense* attention/GDN projections. The modelopt - ``MIXED_PRECISION`` checkpoints tag ``self_attn.{q,k,v,o}_proj`` and - ``linear_attn.{in_proj_qkv,in_proj_z,out_proj}`` with ``quant_algo`` ``FP8`` (fp8-e4m3 - weight + a scalar ``weight_scale``; W8A16). Returns ``"fp8_pertensor"`` when present, - else ``"none"`` (NVFP4 dense weights are covered by ``dense_quant`` / ``lm_head_quant``).""" - get = _quant_accessor(hf_config) - if get is None: - return "none" - layers = get("quantized_layers") or {} - if not isinstance(layers, dict): - return "none" - for name, spec in layers.items(): - algo = str((spec or {}).get("quant_algo", "")).lower() - if algo == "fp8" and (".self_attn." in name or ".linear_attn." in name): - return "fp8_pertensor" - return "none" + # the engine reads this tag for its MoE strategy decisions; every module takes its own scheme from the QuantConfig when it is built + scheme = QuantConfig.from_hf(hf_config).scheme_for_name("model.language_model.layers.0.mlp.experts.0.gate_proj") + if scheme is None: + return "none", None + return str(scheme.kind), scheme.weight.group if scheme.kind is QuantKind.FP8_BLOCK else None def _layer_types(text: Any) -> list[str]: @@ -164,28 +63,7 @@ def parse_config(hf_config: Any) -> ModelConfig: else {k: v for k, v in rope_params.items() if not isinstance(v, (list, dict))} ) - expert_quant, weight_block_size = _fp8_block_quant(hf_config) - if expert_quant == "none": - expert_quant = _expert_quant(hf_config) # nvfp4 / mixed-precision modelopt - # Dense attention/GDN quant is independent of the routed experts (block-fp8 already - # quantizes both, so only probe for per-tensor FP8 when experts aren't block-fp8). - attn_quant = "none" if expert_quant == "fp8_block" else _attn_quant(hf_config) - # NVFP4 checkpoints store the dense MLP projections (shared_expert; dense non-MoE MLP) as - # packed FP4 exactly like the routed experts -- independent of whether attention is FP8 - # (mixed) or bf16 (pure NVFP4). Keep them native FP4 (W4A16) whenever the experts are - # NVFP4. The lm_head is detected separately (only the mixed checkpoint quantizes it). - # MoE-NVFP4 keeps the shared_expert dense MLP native FP4 (expert_quant=="nvfp4"); a dense - # (non-MoE) modelopt checkpoint instead tags the bare .mlp.{gate,up,down}_proj as NVFP4. - dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config) - lm_head_quant = _lm_head_quant(hf_config) - - # compressed-tensors NVFP4 (dense Qwen3.6-27B): the attention (q/k/v/o, GDN out_proj) AND - # the dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms stay bf16. Wire the shared - # W4A16 kernels (attn_quant=="nvfp4" routes the attention/GDN linears through them too). - if _compressed_tensors_nvfp4(hf_config): - attn_quant = "nvfp4" - dense_quant = "nvfp4" - lm_head_quant = "none" + expert_quant, weight_block_size = _expert_quant(hf_config, text) # Dense variants (e.g. Qwen3.6-27B) report num_experts==0: route the decoder MLP through # the dense Qwen3_5DenseMLP instead of the MoE block. @@ -254,9 +132,6 @@ def parse_config(hf_config: Any) -> ModelConfig: attention_groups=groups, expert_quant=expert_quant, weight_block_size=weight_block_size, - attn_quant=attn_quant, - dense_quant=dense_quant, - lm_head_quant=lm_head_quant, ) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index 45adb48b5..d73df9bab 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -1,171 +1,222 @@ +"""Qwen3.5 / 3.6 / 3.8 checkpoint reader. + +The dense pass reads every Linear module under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, block-fp8, ModelOpt and llm-compressor exports in any mix all land as the model's state dict. Routed experts are read by the offload cache (``nvfp4_expert_spec`` / ``iter_expert_pieces``); only bf16 stacked experts and resident block-fp8 experts come from here. +""" + from __future__ import annotations -import json -import os import re from typing import Iterator import safetensors import torch -from freetoken.layers.quantization import QuantKind from freetoken.distributed import get_tp_info from freetoken.kernel.triton.nvfp4_dequant import dequant_nvfp4 -from freetoken.models.loader import ( - CT_SCALE_SUFFIXES, - ShardReader, - ct_bf16_fuse, - ct_nvfp4_fuse, - drop_page_cache, - iter_weight_files, - nvfp4_parts_ct, -) +from freetoken.layers.quantization import QuantConfig, QuantKind, QuantScheme, get_quant_config +from freetoken.models.loader import ShardReader, iter_weight_files from freetoken.models.nvfp4_banks import Nvfp4ExpertSourceSpec -from freetoken.utils import cached_load_hf_config, download_hf_weight +from freetoken.models.register import ModelSpec, get_model_spec +from freetoken.utils import cached_load_hf_config from tqdm import tqdm -from .config import _compressed_tensors_nvfp4, parse_config +from .config import parse_config -# Expert weights are stored pre-fused per layer: experts.gate_up_proj / experts.down_proj. -_PACKED_EXPERT_PATTERN = re.compile( - r"^model\.layers\.\d+\.mlp\.experts\.(gate_up_proj|down_proj)$" -) - -# NVFP4 routed experts (nvidia modelopt checkpoint): per-expert, un-fused, under the raw -# ``model.language_model.layers.N.mlp.experts.E.{proj}`` key. Matched against the RAW -# weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP -# head's ``mtp.layers.N.mlp.experts.*`` tensors (served text-only, dropped). -_NVFP4_EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") -_NVFP4_EXPERT_KEY_RE = re.compile( +# bf16 checkpoints store the routed experts pre-stacked per layer +_STACKED_EXPERT_RE = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate_up_proj|down_proj)$") +# per-expert tensors of a quantized checkpoint: the offload cache's expert reader takes these +_EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") +# the ``model.language_model.`` anchor excludes the MTP head's ``mtp.layers.N.mlp.experts.*`` +_EXPERT_KEY_RE = ( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" -) -_NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( - key_pattern=_NVFP4_EXPERT_KEY_RE, - proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, - layer_to_bank=lambda layer, config: layer, # every layer is MoE - desc="Qwen3.5 NVFP4 experts", + r"(?Pgate_proj|up_proj|down_proj)\.(?P{kinds})$" ) -# Suffixes of the per-tensor modelopt quant scales; consumed alongside their ``.weight``, -# never yielded on their own. -_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +# role -> the expert bank reader's canonical (ModelOpt) tensor kind +_BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"} -# Gemma-style (1+weight) RMSNorm weights. Excludes GDN gated norm (linear_attn.norm), -# which is a standard weight*x norm. +# Gemma-style (1+weight) RMSNorm weights; the GDN gated norm (linear_attn.norm) is a plain weight*x norm _GEMMA_NORM_SUFFIXES = ( ".input_layernorm.weight", ".post_attention_layernorm.weight", ".self_attn.q_norm.weight", ".self_attn.k_norm.weight", ) -# shared-expert gate/up merge -> shared_expert.gate_up_proj -_SHARED_GATE = ".mlp.shared_expert.gate_proj.weight" -_SHARED_UP = ".mlp.shared_expert.up_proj.weight" - -# Fused projections: concat checkpoint matrices in this exact order to match the model's -# LinearColParallelMerged split. fused_suffix -> ordered parts. -_FUSIONS: dict[str, tuple[str, ...]] = { - ".self_attn.qkv_proj.weight": ( - ".self_attn.q_proj.weight", ".self_attn.k_proj.weight", ".self_attn.v_proj.weight", - ), - ".linear_attn.in_proj.weight": ( - ".linear_attn.in_proj_qkv.weight", ".linear_attn.in_proj_z.weight", - ".linear_attn.in_proj_b.weight", ".linear_attn.in_proj_a.weight", - ), - # Dense (non-MoE) layer MLP: merge gate|up -> gate_up_proj. Only fires on a bare - # ``.mlp.gate_proj``; ``.mlp.shared_expert.gate_proj`` (MoE) does not end with this. - ".mlp.gate_up_proj.weight": ( - ".mlp.gate_proj.weight", ".mlp.up_proj.weight", - ), -} - - -def _dequant_fp8_weight(weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor: - """Weight-only FP8 -> bf16 (per-tensor static scale). Activations stay bf16 (W8A16), - which is at least as precise as the checkpoint's intended W8A8.""" - return weight.to(torch.bfloat16) * weight_scale.to(torch.bfloat16) - - -def _dequant_nvfp4_weight( - weight: torch.Tensor, weight_scale: torch.Tensor, weight_scale_2: torch.Tensor -) -> torch.Tensor: - """Dense NVFP4 -> bf16 (W4A16): ``fp4 * block_scale * global_scale``. ``weight`` is - [O, IN//2] uint8, ``weight_scale`` [O, IN//16] fp8-e4m3, ``weight_scale_2`` the per-tensor - global scalar (broadcast to per-row, matching the offload-cache dequant kernel).""" - # The dequant kernel is GPU-only; the checkpoint-conversion path loads dense weights on - # CPU, so run on CUDA and return on the caller's device (no-op when already on GPU). - orig_device = weight.device - if orig_device.type != "cuda": - dev = torch.device("cuda") - weight, weight_scale, weight_scale_2 = ( - weight.to(dev), weight_scale.to(dev), weight_scale_2.to(dev) - ) - out_features = weight.shape[0] - global_scale = weight_scale_2.reshape(1).to(torch.float16).expand(out_features).contiguous() - slots = torch.zeros(1, dtype=torch.int32, device=weight.device) - out = dequant_nvfp4( - weight.unsqueeze(0).contiguous(), - weight_scale.unsqueeze(0).contiguous(), - global_scale.unsqueeze(0), - slots, - dtype=torch.bfloat16, - )[0] - return out.to(orig_device) - - -def _load_maybe_quantized(f, raw_name: str, keyset: set[str]) -> torch.Tensor: - """Load ``raw_name``; if it is a quantized ``.weight`` with sibling modelopt scales in - the same shard, dequantize to bf16 (NVFP4 if ``weight_scale_2`` present, else FP8). - Plain bf16 weights pass through unchanged.""" - tensor = f.get_tensor(raw_name) - if not raw_name.endswith(".weight"): - return tensor - base = raw_name[: -len(".weight")] - if base + ".weight_scale_2" in keyset: # NVFP4 (two-level block scale) - return _dequant_nvfp4_weight( - tensor, f.get_tensor(base + ".weight_scale"), f.get_tensor(base + ".weight_scale_2") - ) - if base + ".weight_scale" in keyset: # FP8 (per-tensor scale) - return _dequant_fp8_weight(tensor, f.get_tensor(base + ".weight_scale")) - return tensor +# leaves the model builds as Linear layers: only their tensors are read under the QuantConfig, the rest passes through as stored +_LINEAR_LEAVES = frozenset({ + "q_proj", "k_proj", "v_proj", "o_proj", "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj", + "gate_proj", "up_proj", "down_proj", "gate", "shared_expert_gate", "lm_head", +}) +# activation scales of modules whose scheme carries no input_scale role +_DROPPED_SUFFIXES = frozenset({"input_scale", "input_global_scale"}) +_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn, "e2m1": torch.uint8} +_QUANT_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2, torch.uint8, torch.int8) def _rename(raw_name: str) -> str | None: - """HF key -> FreeToken state-dict key, or None to skip.""" + """Checkpoint key -> FreeToken state-dict key, or None to skip.""" if raw_name.startswith(("mtp.", "model.visual.", "visual.")): return None - # ModelOpt FP8 KV-cache static scales (full-attention layers only). FreeToken keeps the - # KV cache in the engine's native precision (>= the checkpoint's quantized KV), so these - # per-tensor q/k/v scales are unused -- drop them rather than fail as unexpected keys. + # static KV-cache scales of the quantizers; the KV cache runs in the engine's dtype if raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): return None - name = raw_name - if name.startswith("model.language_model."): - name = "model." + name[len("model.language_model.") :] - elif name.startswith("language_model."): - name = "model." + name[len("language_model.") :] - return name + if raw_name.startswith("model.language_model."): + return "model." + raw_name[len("model.language_model."):] + if raw_name.startswith("language_model."): + return "model." + raw_name[len("language_model."):] + return raw_name def _is_gemma_norm(name: str) -> bool: return name == "model.norm.weight" or name.endswith(_GEMMA_NORM_SUFFIXES) -def _try_fuse( - name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]] -) -> tuple[str, torch.Tensor] | tuple[()] | None: - """buffer a fusion part; return merged ``(name, tensor)`` once all parts arrive, - ``()`` while incomplete, ``None`` if not a fusion part.""" - for fused_suffix, parts in _FUSIONS.items(): - for idx, part in enumerate(parts): - if name.endswith(part): - key = name[: -len(part)] + fused_suffix - slots = buf.setdefault(key, {}) - slots[idx] = tensor - if len(slots) == len(parts): - del buf[key] - return key, torch.cat([slots[i] for i in range(len(parts))], dim=0) - return () - return None +def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: + """Per-tensor scalar or per-channel ``[rows, 1]`` fp8 scale -> fp32 ``[rows]``; any other count is refused rather than broadcast onto the wrong rows.""" + flat = scale.reshape(-1).to(torch.float32) + if flat.numel() == 1: + return flat.expand(rows).contiguous() + if flat.numel() != rows: + raise ValueError( + f"fp8 weight_scale has {flat.numel()} elements for a weight with {rows} output rows " + f"(shape {tuple(scale.shape)}); expected 1 or {rows}" + ) + return flat.contiguous() + + +def _dequant_nvfp4(weight: torch.Tensor, weight_scale: torch.Tensor, weight_global: torch.Tensor) -> torch.Tensor: + """Packed NVFP4 -> bf16 on CUDA (the kernel is GPU-only, the converter reads on CPU), returned on the caller's device.""" + device = weight.device + if device.type != "cuda": + weight, weight_scale, weight_global = (t.to("cuda") for t in (weight, weight_scale, weight_global)) + slots = torch.zeros(1, dtype=torch.int32, device=weight.device) + out = dequant_nvfp4( + weight.unsqueeze(0).contiguous(), weight_scale.unsqueeze(0).contiguous(), weight_global.unsqueeze(0), + slots, dtype=torch.bfloat16, + )[0] + return out.to(device) + + +def _dequant(scheme: QuantScheme, part: dict[str, torch.Tensor]) -> torch.Tensor: + """bf16 weight of a module the checkpoint quantized but the family serves unquantized.""" + weight = part["weight"] + if scheme.kind is QuantKind.FP8_TENSOR: + return (weight.to(torch.float32) * part["weight_scale"][:, None]).to(torch.bfloat16) + if scheme.kind is QuantKind.FP8_BLOCK: + from freetoken.kernel.triton.fp8_block_linear import dequant_block_fp8 + + return dequant_block_fp8(weight, part["weight_scale_inv"]) + if scheme.kind is QuantKind.NVFP4: + return _dequant_nvfp4(weight, part["weight_scale"], part["weight_global"]) + raise NotImplementedError(f"no bf16 dequantization for {scheme}") + + +class _DenseReader: + """Routes each Linear tensor to the buffer its module's scheme declares; packed projections are concatenated per role once every part is in.""" + + def __init__(self, quant: QuantConfig | None, spec: ModelSpec) -> None: + self.quant = quant + self.groups = {fused: parts for fused, parts in spec.packed_modules_mapping if fused != "experts"} + self.by_part: dict[str, list[tuple[str, int]]] = {} + for fused, parts in self.groups.items(): + for idx, part in enumerate(parts): + self.by_part.setdefault(part, []).append((fused, idx)) + # target module -> (part count, {part: {role: tensor}}, {part: the roles its module stores}) + self.pending: dict[str, tuple[int, dict[int, dict[str, torch.Tensor]], dict[int, set[str]]]] = {} + + def scheme(self, module: str) -> QuantScheme | None: + return None if self.quant is None else self.quant.scheme_for(module) + + def stored(self, module: str) -> QuantScheme | None: + """The scheme the checkpoint stores ``module`` under, before the family's unquantized_modules.""" + if self.quant is None: + return None + return self.quant.scheme_for_name(self.quant.name_map.to_checkpoint(module)[0]) + + def target(self, module: str) -> tuple[str, int, int]: + """``(fused module, part index, part count)``; a standalone linear is its own single-part target.""" + parent, _, leaf = module.rpartition(".") + candidates = self.by_part.get(leaf) + if not candidates: + return module, 0, 1 + if len(candidates) > 1: + # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as gdn.py + split = self.scheme(f"{parent}.in_proj_qkvz") is not None + keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"} + candidates = [c for c in candidates if c[0] in keep] + fused, idx = candidates[0] + return f"{parent}.{fused}", idx, len(self.groups[fused]) + + def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None: + """Take one tensor; the emitted ``[(name, tensor)]`` once its target module is complete, ``[]`` before, None if ``name`` is not a Linear's tensor.""" + module, _, suffix = name.rpartition(".") + if module.rpartition(".")[2] not in _LINEAR_LEAVES: + return None + stored = self.stored(module) + roles = {"weight": "weight"} if stored is None else {e.name: r for r, e in self.quant.storage(stored).items()} + role = roles.get(suffix) + if role is None: + if suffix in _DROPPED_SUFFIXES: + return [] + raise ValueError( + f"{name}: the checkpoint's quant config declares {module} {stored or 'unquantized'}, stored as {sorted(roles)}" + ) + if stored is None and tensor.dtype in _QUANT_DTYPES: + raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized") + target, idx, count = self.target(module) + _, parts, expected = self.pending.setdefault(target, (count, {}, {})) + parts.setdefault(idx, {})[role] = tensor + expected[idx] = set(roles.values()) + if len(parts) < count or any(set(parts[i]) != expected[i] for i in parts): + return [] + del self.pending[target] + return self._emit(target, [parts[i] for i in range(count)], stored) + + def _emit(self, target: str, parts: list[dict[str, torch.Tensor]], stored: QuantScheme | None): + if stored is not None: + parts = [self._check(target, stored, part) for part in parts] + if self.scheme(target) is None: + parts = [{"weight": _dequant(stored, part)} for part in parts] + out = [] + for role in parts[0]: + tensors = [part[role] for part in parts] + if role == "input_scale": + # fused parts read the same activation, so ModelOpt calibrates one range for them: max is exact then and safe if they drift + value = torch.stack(tensors).max() + else: + value = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) + out.append((f"{target}.{role}", value)) + return out + + def _check(self, target: str, scheme: QuantScheme, part: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Validate one part against ``scheme`` and put its scales in the layer's form.""" + part = { + role: 1.0 / tensor.to(torch.float32) if self.quant.storage(scheme)[role].reciprocal else tensor + for role, tensor in part.items() + } + weight = part["weight"] + if weight.dtype is not _ELEM_DTYPES[scheme.weight.elem]: + raise ValueError(f"{target}: weight is {weight.dtype} but the checkpoint's quant config declares {scheme}") + rows, cols = weight.shape[0], weight.shape[1] * (2 if scheme.weight.elem == "e2m1" else 1) + block_rows, block_cols = scheme.weight.group or (1, 1) + scale_role = "weight_scale_inv" if "weight_scale_inv" in part else "weight_scale" + out = dict(part) + if block_cols < 0: + out[scale_role] = _per_row_scale(part[scale_role], rows) + else: + if rows % block_rows or cols % block_cols: + raise ValueError(f"{target}: {rows}x{cols} weight is not a multiple of the {block_rows}x{block_cols} scale block of {scheme}") + expected = (rows // block_rows, cols // block_cols) + if tuple(part[scale_role].shape) != expected: + raise ValueError(f"{target}: {scale_role} is {tuple(part[scale_role].shape)}, expected {expected} for {scheme}") + if scheme.weight.scale == "e4m3" and part[scale_role].dtype is not torch.float8_e4m3fn: + raise ValueError(f"{target}: {scale_role} is {part[scale_role].dtype} but {scheme} stores e4m3 scales") + if "weight_global" in part: + g = part["weight_global"].reshape(-1).to(torch.float32) + if g.numel() != 1: + raise ValueError(f"{target}: weight_global has {g.numel()} elements, expected one per-tensor scale") + out["weight_global"] = g.to(torch.float16).expand(rows).contiguous() + if "input_scale" in part: + out["input_scale"] = part["input_scale"].reshape(()).to(torch.float32) + return out def iter_weights( @@ -175,484 +226,45 @@ def iter_weights( include_moe_experts: bool, include_non_moe: bool, ) -> Iterator[tuple[str, torch.Tensor]]: + """Yield the dense weights fused to the model's buffers, and the routed experts only where a resident path takes them from here: bf16 stacked experts as stored, block-fp8 experts restacked per layer. + + Per-expert NVFP4 experts always come from the offload cache's expert reader. + """ + if get_tp_info().size > 1: + raise NotImplementedError("qwen3_5_moe weight loading supports TP=1 only") hf_config = cached_load_hf_config(model_path) config = parse_config(hf_config) - if _compressed_tensors_nvfp4(hf_config): - # Dense compressed-tensors NVFP4 (e.g. Qwen3.6-27B): attn (q/k/v/o, GDN out_proj) + - # dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms bf16. - yield from _iter_weights_compressed_tensors( - model_path, device, - include_non_moe=include_non_moe, include_moe_experts=include_moe_experts, - nvfp4=config.dense_quant == "nvfp4", - ) - return - if config.expert_quant == "fp8_block": - # Dense (attn/GDN/shared-expert) weights are always block-fp8; routed experts are - # yielded here only for the resident path (include_moe_experts=True). Under offload - # they are excluded and loaded from expert pieces instead. - yield from _iter_weights_fp8( - model_path, device, - include_non_moe=include_non_moe, include_moe_experts=include_moe_experts, - ) - return - if config.attn_quant == "fp8_pertensor": - # modelopt MIXED_PRECISION: dense attn/GDN projections kept per-tensor FP8 (fp8 - # weight + per-row scale, W8A16 kernel); NVFP4 dense (shared_expert/lm_head) kept - # native FP4 (W4A16) when dense_quant=="nvfp4", else dequantized to bf16; routed - # NVFP4 experts excluded (offload cache). - yield from _iter_weights_attn_fp8( - model_path, device, - include_non_moe=include_non_moe, include_moe_experts=include_moe_experts, - dense_nvfp4=config.dense_quant == "nvfp4", - lmhead_nvfp4=config.lm_head_quant == "nvfp4", - ) - return - tp_info = get_tp_info() - if tp_info.size > 1: - raise NotImplementedError("qwen3_5_moe weight loading currently supports TP=1 only") - - # Pure-NVFP4 checkpoint (bf16 attn): the dense MLP projections (shared_expert) are still - # stored as packed FP4 -- keep them native (W4A16) when dense_quant=="nvfp4" rather than - # dequantizing to bf16. lm_head here is bf16 (pure NVFP4 doesn't quantize it). - dense_nvfp4 = config.dense_quant == "nvfp4" - lmhead_nvfp4 = config.lm_head_quant == "nvfp4" - shared_buf: dict[str, dict[str, torch.Tensor]] = {} - nvfp4_shared_buf: dict[str, dict[str, tuple]] = {} - fuse_buf: dict[str, dict[int, torch.Tensor]] = {} - - for file in tqdm( - iter_weight_files(model_path), - desc="Loading weights", - disable=not tp_info.is_primary(), - ): - with safetensors.safe_open(file, framework="pt", device=str(device)) as f: - keyset = set(f.keys()) - for raw_name in f.keys(): - # Per-expert NVFP4 tensors go to the offload cache (expert pieces), - # not the dense pass. bf16-base stacked experts (experts.gate_up_proj) have no - # ``.mlp.experts..`` so they are unaffected and still hit _PACKED_EXPERT. - if _NVFP4_EXPERT_RE.search(raw_name): - continue - # Standalone modelopt scales are consumed with their .weight, never yielded. - if raw_name.endswith(_SCALE_SUFFIXES): - continue - - name = _rename(raw_name) - if name is None: - continue - - is_expert = _PACKED_EXPERT_PATTERN.match(name) is not None - if is_expert and not include_moe_experts: - continue - if not is_expert and not include_non_moe: - continue - - # NVFP4 dense projections kept native (W4A16) where the model expects them - # (shared_expert); everything else dequantizes to bf16 below as before. - if (dense_nvfp4 or lmhead_nvfp4) and name.endswith(".weight") \ - and raw_name[: -len(".weight")] + ".weight_scale_2" in keyset: - emit = _dense_nvfp4_emit( - f, name[: -len(".weight")], raw_name[: -len(".weight")], - shared_nvfp4=dense_nvfp4, lmhead_nvfp4=lmhead_nvfp4, - shared_buf=nvfp4_shared_buf, - ) - if emit is not _NOT_DENSE_NVFP4: - yield from emit - continue - - tensor = _load_maybe_quantized(f, raw_name, keyset) - - # merge shared-expert gate/up -> gate_up_proj - if name.endswith(_SHARED_GATE) or name.endswith(_SHARED_UP): - prefix = name.rsplit(".mlp.shared_expert.", 1)[0] - slots = shared_buf.setdefault(prefix, {}) - slots["gate" if name.endswith(_SHARED_GATE) else "up"] = tensor - if "gate" in slots and "up" in slots: - merged = torch.cat([slots["gate"], slots["up"]], dim=0) - del shared_buf[prefix] - yield f"{prefix}.mlp.shared_expert.gate_up_proj.weight", merged - continue - - # fuse q/k/v -> qkv_proj and GDN in_proj_{qkv,z,b,a} -> in_proj - fused = _try_fuse(name, tensor, fuse_buf) - if fused is not None: - if fused != (): # () means buffered, not yet complete - yield fused - continue - - if _is_gemma_norm(name): - tensor = tensor + 1.0 # (1 + weight) baked into the stored weight - - yield name, tensor - - assert not shared_buf, f"Incomplete shared-expert merges: {list(shared_buf.keys())}" - assert not nvfp4_shared_buf, f"Incomplete NVFP4 shared-expert merges: {list(nvfp4_shared_buf.keys())}" - assert not fuse_buf, f"Incomplete projection fusions: {list(fuse_buf.keys())}" - - -# ====================================================================================== -# Mixed-precision modelopt checkpoint (per-tensor FP8 attn/GDN + NVFP4 experts/shared/lm_head) -# ====================================================================================== -# FP8 projections fused along the output dim, each part carrying its own scalar weight_scale -# -> a per-output-row scale vector. Keys are the model buffer base (sans .weight/.weight_scale). -_PT_FP8_FUSE: dict[str, tuple[str, ...]] = { - ".self_attn.qkv_proj": ( - ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", - ), - ".linear_attn.in_proj_qkvz": ( - ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", - ), -} -# bf16 (unquantized) GDN b|a projections fused -> in_proj_ba (matches the fp8 split). -_PT_BF16_FUSE: dict[str, tuple[str, ...]] = { - ".linear_attn.in_proj_ba": (".linear_attn.in_proj_b", ".linear_attn.in_proj_a"), -} - - -def _per_row_scale(scalar: torch.Tensor, rows: int) -> torch.Tensor: - """Per-tensor scalar -> per-output-row fp32 vector ``[rows]`` (exact broadcast).""" - return scalar.reshape(1).to(torch.float32).expand(rows) - - -def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor, - act_scale: torch.Tensor | None, buf: dict): - """Buffer an fp8 fusion part ``(weight, scalar, act_scale)``; once all parts arrive emit - the concatenated ``(.weight fp8, .weight_scale per-row fp32)`` plus the shared - ``.input_scale``. ``[]`` while incomplete, ``None`` if ``base`` is not an fp8 fusion part. - - The fused parts all read the *same* activation, so modelopt calibrates one activation - range for all of them and their ``input_scale`` values come out bit-identical (verified on - Qwen3.8-27B-NVFP4: q/k/v all 0.2053571492, GDN qkv/z both 0.1121651828). Taking the max is - therefore exact here, and stays correct if a future checkpoint lets them drift.""" - for fused_suffix, parts in _PT_FP8_FUSE.items(): - for idx, part in enumerate(parts): - if base.endswith(part): - key = base[: -len(part)] + fused_suffix - slots = buf.setdefault(key, {}) - slots[idx] = (weight, scalar, act_scale) - if len(slots) < len(parts): - return [] - del buf[key] - ws = [slots[i][0] for i in range(len(parts))] - ss = [_per_row_scale(slots[i][1], slots[i][0].shape[0]) for i in range(len(parts))] - emit = [ - (key + ".weight", torch.cat(ws, dim=0)), - (key + ".weight_scale", torch.cat(ss, dim=0).contiguous()), - ] - acts = [slots[i][2] for i in range(len(parts))] - if all(a is not None for a in acts): - emit.append((key + ".input_scale", torch.stack( - [a.reshape(()).to(torch.float32) for a in acts]).max())) - return emit - return None - - -# Native-NVFP4 dense projections (W4A16): shared-expert gate/up merged on the output dim, -# down + lm_head standalone. Each carries weight (uint8 [O,IN//2]), block scale (fp8 -# [O,IN//16]), and a per-output-row global scale (weight_scale_2 broadcast, fp16 [O]). The -# fused gate|up concatenates all three (each part keeps its own global), so it is exact. -_SHARED_GATE_BASE = ".mlp.shared_expert.gate_proj" -_SHARED_UP_BASE = ".mlp.shared_expert.up_proj" -# The MoE shared-expert MLP and the dense (non-MoE) decoder MLP have identical native-FP4 -# structure (gate|up merged -> gate_up_proj + standalone down_proj); they differ only in the -# ``.mlp.shared_expert.`` vs bare ``.mlp.`` infix. ``endswith(".mlp.gate_proj")`` is False for -# ``.mlp.shared_expert.gate_proj``, and routed experts are excluded upstream (_NVFP4_EXPERT_RE). -_NVFP4_MLP_LAYOUTS = ( - (".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj", - ".mlp.shared_expert.down_proj", ".mlp.shared_expert."), - (".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj", ".mlp."), -) + stacked = include_moe_experts and config.is_moe and config.expert_quant == "none" + if include_non_moe or stacked: + reader = _DenseReader(get_quant_config(), get_model_spec(hf_config.architectures[0])) if include_non_moe else None + yield from _iter_shards(model_path, device, reader, stacked=stacked) + if include_moe_experts and config.is_moe and config.expert_quant == "fp8_block": + yield from _resident_fp8_experts(model_path, config) -def _nvfp4_parts(f, raw_base: str): - """Load a native NVFP4 weight as ``(packed uint8 [O, IN//2], block scale fp8 [O, IN//16], - per-output-row global fp16 [O])`` -- the dense W4A16 kernels' expected buffers.""" - w = f.get_tensor(raw_base + ".weight") # uint8 packed FP4 (2 codes/byte) - s = f.get_tensor(raw_base + ".weight_scale") # fp8-e4m3 per-16 block scale - g2 = f.get_tensor(raw_base + ".weight_scale_2") # per-tensor global scalar - g = g2.reshape(1).to(torch.float16).expand(w.shape[0]).contiguous() - return w, s, g - - -# Sentinel: ``base`` is not a dense projection the model keeps native NVFP4 (caller dequantizes). -_NOT_DENSE_NVFP4 = object() - - -def _dense_nvfp4_emit( - f, base: str, raw_base: str, *, shared_nvfp4: bool, lmhead_nvfp4: bool, shared_buf: dict -): - """For a dense ``.weight`` whose checkpoint has a ``weight_scale_2`` (NVFP4), return the list - of ``(key, tensor)`` to yield as native FP4 -- ``(.weight uint8, .weight_scale fp8 block, - .weight_global fp16 per-row)`` -- when the model keeps that layer native: - - * the MoE ``shared_expert.{gate,up,down}_proj`` OR the dense (non-MoE) ``.mlp.{gate,up,down} - _proj`` when ``shared_nvfp4`` (gate/up merged -> ``gate_up_proj``, each part keeping its own - global scale so the fused weight is exact); - * ``lm_head`` when ``lmhead_nvfp4``. - - Returns ``[]`` while a gate/up merge is still buffered, or ``_NOT_DENSE_NVFP4`` if the model - does not keep this layer native (the caller dequantizes to bf16 exactly as before). Shared by - the mixed-FP8 dense pass and the default (pure-NVFP4) dense pass.""" - is_lmhead = base == "lm_head" or base.endswith(".lm_head") - if lmhead_nvfp4 and is_lmhead: - w, s, g = _nvfp4_parts(f, raw_base) - return [(base + ".weight", w), (base + ".weight_scale", s), (base + ".weight_global", g)] - if not shared_nvfp4: - return _NOT_DENSE_NVFP4 - for gate_b, up_b, down_b, infix in _NVFP4_MLP_LAYOUTS: - if base.endswith(down_b): - w, s, g = _nvfp4_parts(f, raw_base) - return [(base + ".weight", w), (base + ".weight_scale", s), (base + ".weight_global", g)] - if base.endswith(gate_b) or base.endswith(up_b): - w, s, g = _nvfp4_parts(f, raw_base) - prefix = base.rsplit(infix, 1)[0] + infix - slots = shared_buf.setdefault(prefix, {}) - slots["gate" if base.endswith(gate_b) else "up"] = (w, s, g) - if "gate" not in slots or "up" not in slots: - return [] - gw, gs, gg = slots["gate"] - uw, us, ug = slots["up"] - del shared_buf[prefix] - pre = f"{prefix}gate_up_proj" - return [ - (pre + ".weight", torch.cat([gw, uw], dim=0)), - (pre + ".weight_scale", torch.cat([gs, us], dim=0)), - (pre + ".weight_global", torch.cat([gg, ug], dim=0)), - ] - return _NOT_DENSE_NVFP4 - - -def _iter_weights_attn_fp8( - model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool, - dense_nvfp4: bool = False, lmhead_nvfp4: bool = False, -) -> Iterator[tuple[str, torch.Tensor]]: - """Dense pass for the modelopt MIXED_PRECISION Qwen3.5 checkpoint. - - Per-tensor FP8 attn/GDN projections (``self_attn.{q,k,v,o}_proj``, ``linear_attn. - {in_proj_qkv,in_proj_z,out_proj}``) are kept fp8-e4m3 and yielded as ``.weight`` (fp8) + - ``.weight_scale`` (per-output-row fp32) instead of dequantized to bf16 -- this halves the - decode weight traffic of the dense backbone. q/k/v -> ``qkv_proj``, GDN qkv|z -> - ``in_proj_qkvz`` (fp8) and b|a -> ``in_proj_ba`` (bf16). NVFP4 dense weights - (shared_expert, lm_head): kept native FP4 -- ``.weight`` (uint8) + ``.weight_scale`` - (fp8 block) + ``.weight_global`` (fp16 per-row) for the W4A16 kernels -- when - ``dense_nvfp4`` else dequantized to bf16. Routed NVFP4 experts are excluded (served by - the offload cache). Gemma (1+w) norms get +1.""" - if get_tp_info().size > 1: - raise NotImplementedError("qwen3_5_moe weight loading currently supports TP=1 only") - if not include_non_moe: - return # experts-only call: NVFP4 experts are loaded by the offload bank provider - - tp_info = get_tp_info() - fp8_buf: dict[str, dict[int, tuple]] = {} - bf16_buf: dict[str, dict[int, torch.Tensor]] = {} - shared_buf: dict[str, dict[str, torch.Tensor]] = {} - nvfp4_shared_buf: dict[str, dict[str, tuple]] = {} - - for file in tqdm( - iter_weight_files(model_path), - desc="Loading mixed-fp8 weights", - disable=not tp_info.is_primary(), - ): +def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | None, *, stacked: bool): + for file in tqdm(iter_weight_files(model_path), desc="Loading weights", disable=not get_tp_info().is_primary()): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: - keyset = set(f.keys()) for raw_name in f.keys(): - if _NVFP4_EXPERT_RE.search(raw_name): - continue # routed experts -> offload cache - if raw_name.endswith(_SCALE_SUFFIXES): - continue # scales consumed with their .weight - - name = _rename(raw_name) - if name is None: - continue - if _PACKED_EXPERT_PATTERN.match(name) is not None: - continue # no packed experts in this checkpoint; guard anyway - - if name.endswith(".weight"): - base = name[: -len(".weight")] - raw_base = raw_name[: -len(".weight")] - has_s2 = raw_base + ".weight_scale_2" in keyset - has_s = raw_base + ".weight_scale" in keyset - if has_s and not has_s2: # per-tensor FP8 dense projection - w = f.get_tensor(raw_name) # fp8-e4m3, kept verbatim - sc = f.get_tensor(raw_base + ".weight_scale") - # modelopt's calibrated activation scale: kept (not dropped with the - # other scale suffixes) so batched decode can run W8A8 instead of - # W8A16. Absent -> the layer stays on the W8A16 kernel. - act = (f.get_tensor(raw_base + ".input_scale") - if raw_base + ".input_scale" in keyset else None) - emit = _pt_fp8_fuse(base, w, sc, act, fp8_buf) - if emit is not None: - yield from emit - continue - # standalone fp8 (self_attn.o_proj, linear_attn.out_proj) - yield base + ".weight", w - yield base + ".weight_scale", _per_row_scale(sc, w.shape[0]).contiguous() - if act is not None: - yield base + ".input_scale", act.reshape(()).to(torch.float32) - continue - if has_s2: # NVFP4 dense: keep native (W4A16) where the model expects it - emit = _dense_nvfp4_emit( - f, base, raw_base, shared_nvfp4=dense_nvfp4, - lmhead_nvfp4=lmhead_nvfp4, shared_buf=nvfp4_shared_buf, - ) - if emit is not _NOT_DENSE_NVFP4: - yield from emit - continue - # NVFP4 -> bf16 (shared_expert, lm_head; dense_nvfp4 off); plain bf16 passes through. - tensor = _load_maybe_quantized(f, raw_name, keyset) - emit = _ct_bf16_fuse(base, tensor, bf16_buf, _PT_BF16_FUSE) - if emit is not None: - yield from emit - continue - else: - tensor = f.get_tensor(raw_name) - - # shared-expert gate/up -> gate_up_proj (bf16, dequantized above) - if name.endswith(_SHARED_GATE) or name.endswith(_SHARED_UP): - prefix = name.rsplit(".mlp.shared_expert.", 1)[0] - slots = shared_buf.setdefault(prefix, {}) - slots["gate" if name.endswith(_SHARED_GATE) else "up"] = tensor - if "gate" in slots and "up" in slots: - merged = torch.cat([slots["gate"], slots["up"]], dim=0) - del shared_buf[prefix] - yield f"{prefix}.mlp.shared_expert.gate_up_proj.weight", merged - continue - - if _is_gemma_norm(name): - tensor = tensor + 1.0 # (1 + weight) baked into the stored norm weight - - yield name, tensor - - assert not fp8_buf, f"Incomplete fp8 fusions: {list(fp8_buf.keys())}" - assert not bf16_buf, f"Incomplete bf16 fusions: {list(bf16_buf.keys())}" - assert not shared_buf, f"Incomplete shared-expert merges: {list(shared_buf.keys())}" - assert not nvfp4_shared_buf, f"Incomplete NVFP4 shared-expert merges: {list(nvfp4_shared_buf.keys())}" - - -# ====================================================================================== -# compressed-tensors NVFP4 checkpoint (dense Qwen3.x, e.g. Qwen3.6-27B) -# ====================================================================================== -# NVFP4 (W4A16) targets every Linear except the per-module ``ignore`` list (lm_head, GDN -# in_proj_*, vision, mtp). Storage differs from modelopt: ``weight_packed`` (uint8 [O, IN//2]) -# + ``weight_scale`` (fp8-e4m3 block [O, IN//16]) + a scalar ``weight_global_scale``. The -# stored global is the *quant-side* scale, so the dequant/native global is its reciprocal -# (``1/weight_global_scale``) -- vLLM inverts it identically. Dense MLP gate/up and attention -# q/k/v fuse on the output dim into ``gate_up_proj`` / ``qkv_proj`` (each part keeps its own -# global, so the fused FP4 weight is exact). GDN ``in_proj_{qkv,z,b,a}`` stay bf16 -> ``in_proj``. -_CT_NVFP4_FUSE: dict[str, tuple[str, ...]] = { - ".self_attn.qkv_proj": (".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj"), - ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj"), -} -_CT_BF16_FUSE: dict[str, tuple[str, ...]] = { - ".linear_attn.in_proj": ( - ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", - ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", - ), -} -# The scale suffixes and parts/fuse machinery are shared with muse_glimmer and live -# in models/loader.py. -_CT_SCALE_SUFFIXES = CT_SCALE_SUFFIXES -_nvfp4_parts_ct = nvfp4_parts_ct -_ct_bf16_fuse = ct_bf16_fuse - - -def _ct_nvfp4_fuse(base: str, parts_tuple: tuple, buf: dict): - return ct_nvfp4_fuse(base, parts_tuple, buf, _CT_NVFP4_FUSE) - - -def _iter_weights_compressed_tensors( - model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool, - nvfp4: bool, -) -> Iterator[tuple[str, torch.Tensor]]: - """Dense pass for a compressed-tensors NVFP4 checkpoint (e.g. Qwen3.6-27B). - - Keeps the NVFP4 attention (q/k/v/o, GDN out_proj) and dense MLP (gate/up/down) native - (W4A16) -- ``.weight`` (uint8) + ``.weight_scale`` (fp8 block) + ``.weight_global`` (fp16 - per-row) -- when ``nvfp4``; otherwise dequantizes each to bf16. q/k/v -> ``qkv_proj``, dense gate/up -> ``gate_up_proj`` (output-dim concat). - GDN ``in_proj_{qkv,z,b,a}`` stay bf16 -> fused ``in_proj``; ``conv1d``/``A_log``/``dt_bias``/ - gated ``norm`` pass through (fp32 for A_log/dt_bias). Gemma (1+w) norms get +1. lm_head and - embeddings are bf16. The model is dense (no routed experts), so there is no experts pass.""" - if get_tp_info().size > 1: - raise NotImplementedError("qwen3_5_moe weight loading currently supports TP=1 only") - if not include_non_moe: - return # dense checkpoint: no routed experts to load - - tp_info = get_tp_info() - nvfp4_buf: dict[str, dict[int, tuple]] = {} - bf16_buf: dict[str, dict[int, torch.Tensor]] = {} - - def _emit_bf16_weight(name: str, tensor: torch.Tensor): - """Plain bf16 ``.weight``: GDN in_proj fusion, Gemma (1+w) norms, else passthrough.""" - base = name[: -len(".weight")] - emit = _ct_bf16_fuse(base, tensor, bf16_buf, _CT_BF16_FUSE) - if emit is not None: - yield from emit - return - if _is_gemma_norm(name): - tensor = tensor + 1.0 # (1 + weight) baked into the stored norm weight - yield name, tensor - - # Scale lookups go through the shard-map reader: a weight_packed's quant scales - # can land in a different shard than the packed weight (the Muse-Glimmer layer-49 - # shape; nothing prevents an llm-compressor Qwen export from splitting the same way). - reader = ShardReader(model_path, device) - try: - for file in tqdm( - reader.files(), - desc="Loading compressed-tensors weights", - disable=not tp_info.is_primary(), - ): - for raw_name in reader.names_in(file): - if raw_name.startswith(("mtp.", "model.visual.", "visual.")): - continue - if raw_name.endswith(_CT_SCALE_SUFFIXES): - continue # consumed with weight_packed (or unused W4A4 activation scales) - name = _rename(raw_name) - if name is None: + if name is None or _EXPERT_RE.search(name): continue - - if raw_name.endswith(".weight_packed"): # NVFP4 projection - base = name[: -len(".weight_packed")] - raw_base = raw_name[: -len(".weight_packed")] - w, s, g = _nvfp4_parts_ct(reader, raw_base) - # GDN in_proj_* compute in bf16 (model contract) but some checkpoints - # (e.g. sakamakismile/Qwen3.6-27B-NVFP4) quantize them too: dequant to - # bf16 here and let the bf16 fusion assemble ``in_proj`` as usual. - if any(base.endswith(p) for ps in _CT_BF16_FUSE.values() for p in ps): - bf16 = _dequant_nvfp4_weight(w, s, g[:1]) - yield from _emit_bf16_weight(base + ".weight", bf16) - continue - if nvfp4: # keep native (W4A16) - emit = _ct_nvfp4_fuse(base, (w, s, g), nvfp4_buf) - if emit is not None: - yield from emit - else: # standalone: o_proj, linear_attn.out_proj, mlp.down_proj - yield base + ".weight", w - yield base + ".weight_scale", s - yield base + ".weight_global", g - continue - # bf16 A-B: dequant FP4 -> bf16, then merge q/k/v + gate/up as bf16. ``g`` is - # already the dequant global (1/weight_global_scale) per row; pass one element. - bf16 = _dequant_nvfp4_weight(w, s, g[:1]) - emit = _ct_bf16_fuse(base, bf16, bf16_buf, _CT_NVFP4_FUSE) - if emit is not None: - yield from emit - else: - yield base + ".weight", bf16 + if _STACKED_EXPERT_RE.match(name): + if stacked: + yield name, f.get_tensor(raw_name) continue - - if name.endswith(".weight"): - yield from _emit_bf16_weight(name, reader.get_tensor(raw_name)) + if reader is None: continue - - # A_log / dt_bias (kept fp32 by the model; the load downcast exempts them). - yield name, reader.get_tensor(raw_name) - finally: - reader.close() - - assert not nvfp4_buf, f"Incomplete NVFP4 fusions: {list(nvfp4_buf.keys())}" - assert not bf16_buf, f"Incomplete bf16 fusions: {list(bf16_buf.keys())}" + tensor = f.get_tensor(raw_name) + emitted = reader.add(name, tensor) + if emitted is not None: + yield from emitted + elif _is_gemma_norm(name): + yield name, tensor + 1.0 # (1 + weight) baked into the stored norm weight + else: + yield name, tensor + if reader is not None and reader.pending: + raise ValueError(f"checkpoint is missing tensors of {sorted(reader.pending)}") def iter_weights_parallel( @@ -677,7 +289,7 @@ def iter_weights_parallel( def _is_expert(raw_name: str) -> bool: name = _rename(raw_name) - return name is not None and _PACKED_EXPERT_PATTERN.match(name) is not None + return name is not None and _STACKED_EXPERT_RE.match(name) is not None for raw_name, tensor in iter_expert_tensors_parallel( model_path, _is_expert, workers=workers, chunk=chunk @@ -686,32 +298,8 @@ def _is_expert(raw_name: str) -> bool: # ====================================================================================== -# Block-FP8 checkpoint (Qwen3.5-35B-A3B-FP8): dense weights + offload expert banks. +# Block-FP8 routed experts (Qwen3.5-35B-A3B-FP8): offload expert pieces and resident stacks. # ====================================================================================== -# fused model buffer suffix -> ordered checkpoint part suffixes (matched without the -# trailing .weight / .weight_scale_inv). Both kinds ride the same fusion (concatenated -# along dim 0); in_proj_ba carries only .weight (b/a stay bf16, no block scale). -_FP8_FUSIONS: dict[str, tuple[str, ...]] = { - ".self_attn.qkv_proj": ( - ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", - ), - ".linear_attn.in_proj_qkvz": ( - ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", - ), - ".linear_attn.in_proj_ba": ( - ".linear_attn.in_proj_b", ".linear_attn.in_proj_a", - ), - ".mlp.shared_expert.gate_up_proj": ( - ".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj", - ), - # Dense (non-MoE) layer MLP: merge gate|up -> gate_up_proj for both the fp8 ``.weight`` - # and the bf16 ``.weight_scale_inv`` (fused per kind by _split_kind). Only a bare - # ``.mlp.gate_proj`` matches; the shared_expert entry above keeps the MoE case. - ".mlp.gate_up_proj": ( - ".mlp.gate_proj", ".mlp.up_proj", - ), -} -_FP8_KIND_SUFFIXES = (".weight_scale_inv", ".weight") # Routed-expert checkpoint key (per-expert, un-fused). ``mtp.layers...`` is excluded by the # ``model.language_model.`` anchor, so the parallel reader only sees the real experts. @@ -721,77 +309,6 @@ def _is_expert(raw_name: str) -> bool: ) -def _split_kind(name: str) -> tuple[str, str]: - """``name`` -> ``(base, kind_suffix)``; ``kind_suffix`` is "" for keys without a - weight/scale suffix (A_log, dt_bias).""" - for suf in _FP8_KIND_SUFFIXES: - if name.endswith(suf): - return name[: -len(suf)], suf - return name, "" - - -def _fp8_fuse(base: str, suf: str, tensor: torch.Tensor, buf: dict) -> tuple[str, torch.Tensor] | tuple[()] | None: - """Buffer a fusion part keyed by (fused_full_name, kind); return the concatenated - ``(name, tensor)`` once all parts for that kind arrive, ``()`` while incomplete, - ``None`` if ``base`` is not a fusion part.""" - for fused_suffix, parts in _FP8_FUSIONS.items(): - for idx, part in enumerate(parts): - if base.endswith(part): - fused_base = base[: -len(part)] + fused_suffix - key = (fused_base + suf, suf) - slots = buf.setdefault(key, {}) - slots[idx] = tensor - if len(slots) == len(parts): - del buf[key] - return fused_base + suf, torch.cat([slots[i] for i in range(len(parts))], dim=0) - return () - return None - - -def _iter_weights_fp8( - model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool = False -) -> Iterator[tuple[str, torch.Tensor]]: - """Yield the block-fp8 weights, renamed + fused to the model buffers. - - fp8 weights (e4m3) and their bf16 ``weight_scale_inv`` pass through verbatim (no dtype - cast -- the engine's load-time cast is a no-op against the fp8/bf16 model buffers). - q/k/v -> qkv_proj, GDN in_proj_qkv|z -> in_proj_qkvz (fp8) and in_proj_b|a -> in_proj_ba - (bf16), shared_expert gate|up -> gate_up_proj; Gemma (1+w) norms get +1. - - Routed experts: skipped under offload (loaded from expert pieces). Under the - resident (non-offload) path ``include_moe_experts`` is True -> per-layer stacked fp8 - experts for the Fp8ResidentMoE buffers are yielded too.""" - if get_tp_info().size > 1: - raise NotImplementedError("qwen3_5_moe fp8 weight loading supports TP=1 only") - if include_non_moe: - fuse_buf: dict = {} - for file in tqdm(iter_weight_files(model_path), desc="Loading fp8 weights", - disable=not get_tp_info().is_primary()): - with safetensors.safe_open(file, framework="pt", device=str(device)) as f: - for raw_name in f.keys(): - name = _rename(raw_name) - if name is None or ".mlp.experts." in name: - continue # routed experts handled below / by the offload cache - tensor = f.get_tensor(raw_name) - base, suf = _split_kind(name) - fused = _fp8_fuse(base, suf, tensor, fuse_buf) - if fused is not None: - if fused != (): - yield fused - continue - if _is_gemma_norm(name): - tensor = tensor + 1.0 # (1 + weight) baked into the stored norm weight - yield name, tensor - assert not fuse_buf, f"Incomplete fp8 fusions: {sorted(k for k, _ in fuse_buf)}" - - if include_moe_experts: - # resident experts: stack the per-expert pieces into the per-layer tensors the resident MoE method declares (pageable host; the engine copies to GPU) - config = parse_config(cached_load_hf_config(model_path)) - if not config.is_moe: - return # dense checkpoint: no routed experts to build as resident banks - yield from _resident_fp8_experts(model_path, config) - - def _resident_fp8_experts(model_path, config): from freetoken.kernel.triton.fp8_block_linear import FP8 @@ -805,7 +322,7 @@ def _resident_fp8_experts(model_path, config): } layers: dict[int, dict[str, torch.Tensor]] = {} placed = [0] * L - for li, e0, e1, piece in iter_expert_pieces(model_path, config, "fp8_block", parallel=None): + for li, e0, e1, piece in iter_expert_pieces(model_path, config, QuantKind.FP8_BLOCK, parallel=None): stack = layers.setdefault(li, {n: torch.empty(shape, dtype=dt) for n, (shape, dt) in shapes.items()}) stack["gate_up_proj"][e0:e1, :I] = piece["gate"] stack["gate_up_proj"][e0:e1, I:] = piece["up"] @@ -821,34 +338,6 @@ def _resident_fp8_experts(model_path, config): assert not layers, f"incomplete resident fp8 experts for layers {sorted(layers)}" -class _ShardReader: - """Opens safetensors shards on demand and serves tensors by name on ``device``.""" - - def __init__(self, folder: str, weight_map: dict, device: torch.device): - self._folder = folder - self._map = weight_map - self._device = device - self._handles: dict = {} - - def get(self, name: str) -> torch.Tensor: - shard = self._map[name] - h = self._handles.get(shard) - if h is None: - h = safetensors.safe_open( - os.path.join(self._folder, shard), framework="pt", device=str(self._device) - ).__enter__() - self._handles[shard] = h - return h.get_tensor(name) - - def close(self) -> None: - for h in self._handles.values(): - try: - h.__exit__(None, None, None) - except Exception: - pass - self._handles.clear() - - def _moe_dims(model_config): L = model_config.num_moe_layers return ( @@ -857,13 +346,6 @@ def _moe_dims(model_config): ) -def _expert_reader(model_path, device): - folder = download_hf_weight(model_path) - with open(os.path.join(folder, "model.safetensors.index.json")) as fh: - weight_map = json.load(fh)["weight_map"] - return _ShardReader(folder, weight_map, device) - - def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): """Block-fp8 routed experts, one piece per expert: ``{gate, up, down}`` fp8 codes and their ``_scale`` (bf16 block ``weight_scale_inv``) companions. Other expert kinds use the generic readers.""" @@ -895,7 +377,7 @@ def locate(raw_name: str): return per_expert_pieces(tensors, locate, tensors_per_expert=6) def _serial(): - reader = _expert_reader(model_path, torch.device("cpu")) + reader = ShardReader(model_path, torch.device("cpu")) try: for li in tqdm(range(L), desc="Loading fp8 experts (serial)", disable=not get_tp_info().is_primary()): for e in range(E): @@ -903,15 +385,26 @@ def _serial(): for proj in ("gate", "up", "down"): for kind, suf in suffix.items(): name = f"{base}.{proj}_proj.{kind}" - yield name, reader.get(name) + yield name, reader.get_tensor(name) finally: reader.close() return per_expert_pieces(_serial(), locate, tensors_per_expert=6) -def nvfp4_expert_spec(model_path: str, config): - return _NVFP4_SOURCE_SPEC +def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: + """The per-expert NVFP4 layout under the checkpoint's dialect names (ModelOpt or llm-compressor).""" + quant = get_quant_config() + stored = quant.stored_tensors(QuantKind.NVFP4) + kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()} + return Nvfp4ExpertSourceSpec( + key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc=f"Qwen3.5 NVFP4 experts ({quant.dialect})", + kind_map=kind_map, + global_reciprocal=stored["weight_global"].reciprocal, + ) __all__ = [ diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 1be3069ae..957c0fa99 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -39,6 +39,8 @@ class ModelSpec: ("in_proj_ba", ("in_proj_b", "in_proj_a")), ("in_proj", ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a")), ) + _EXPERTS_PACKED +# routers the family builds without a quant config (qwen3_5_moe/moe.py), so a checkpoint that quantized them is dequantized at load +_QWEN3_5_UNQUANTIZED = ("*.mlp.gate", "*.mlp.shared_expert_gate") # Qwen3.8's per-layer hyper-connections fuse the down projection with the block-inject rows. _QWEN4_EXP_PACKED = _QWEN3_5_PACKED + ( ("input_mix_weight_down_block_inject", ("input_mix_weight_down", "block_inject_weight")), @@ -114,6 +116,7 @@ class ModelSpec: "Qwen3_5MoEForCausalLM", checkpoint_roots=_LANGUAGE_MODEL_ROOT, packed_modules_mapping=_QWEN3_5_PACKED, + unquantized_modules=_QWEN3_5_UNQUANTIZED, ), # Qwen3.8-Flash-Next (model_type qwen4_exp): multimodal wrapper config (text tower in # text_config, weights under model.language_model.); served text-only. 36 GDN + 12 QSA @@ -133,6 +136,7 @@ class ModelSpec: "Qwen3_5MoEForCausalLM", checkpoint_roots=_LANGUAGE_MODEL_ROOT, packed_modules_mapping=_QWEN3_5_PACKED, + unquantized_modules=_QWEN3_5_UNQUANTIZED, ), # Muse-Glimmer-30B (model_type muse_glimmer): multimodal wrapper config (text tower in # text_config, weights under model.language_model.); served text-only. Dense gated GQA diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 03b11a1de..0cff7eecf 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -87,22 +87,11 @@ def iter_expert_tensors_parallel( Peak host memory is ~(prefetch+1) shards + the banks the caller fills. Order is shard-then-header order (NOT global), so the consumer must place by ``name``. """ + from freetoken.models.loader import safetensors_weight_map from freetoken.utils.hf import download_hf_weight model_path = download_hf_weight(model_path) # resolve hub id -> local (parity w/ serial) - index = os.path.join(model_path, "model.safetensors.index.json") - if os.path.exists(index): - with open(index) as f: - weight_map = json.load(f)["weight_map"] - else: # single-file / no-index checkpoint: map name -> shard from each shard's header - weight_map = {} - for shard in sorted(os.path.basename(p) for p in glob.glob(os.path.join(model_path, "*.safetensors"))): - with open(os.path.join(model_path, shard), "rb") as fh: - n = struct.unpack("= scheme.roles, f"{cls.__name__} does not name every tensor of {scheme}" + + +def test_compressed_tensors_target_classes_other_than_linear_fail_closed(): + q = {"quant_method": "compressed-tensors", "config_groups": {"group_0": { + "targets": ["Embedding"], "weights": {"num_bits": 8, "type": "float", "strategy": "channel"}, "input_activations": {"dynamic": True}}}} + with pytest.raises(NotImplementedError, match="Embedding"): + QuantConfig.from_hf(SimpleNamespace(quantization_config=q)) + + def test_unsupported_dialects_fail_closed(tmp_path): hf_config = SimpleNamespace(architectures=["Qwen3MoeForCausalLM"], quantization_config={"quant_method": "gptq", "bits": 4}) with pytest.raises(NotImplementedError, match="gptq"): diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py new file mode 100644 index 000000000..98f21649c --- /dev/null +++ b/tests/models/test_qwen3_5_moe_weight.py @@ -0,0 +1,553 @@ +"""qwen3_5_moe weight loading against synthetic checkpoints shaped like the released ones. + +Tiny tensors, real key names, dtypes and quantization_config blocks. Each layout is read by +``iter_weights`` and compared with the state dict of the model the engine builds from the same +config; the dense pass has to fill exactly those buffers whatever the checkpoint quantized. +""" + +from __future__ import annotations + +import json +import re + +import pytest +import torch +from safetensors.torch import save_file + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.layers.quantization import set_quant_config +from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces +from freetoken.models.qwen3_5_moe.config import parse_config +from freetoken.models.qwen3_5_moe.weight import iter_weights, nvfp4_expert_spec +from freetoken.models.register import checkpoint_quant_config, get_model_spec +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 +KH, VH, HD = 2, 4, 64 # GDN key / value heads, head dim: qkv rows 512, z rows 256, b|a rows 4 each +QH, KVH, AHD = 4, 2, 64 # attention q / kv heads, head dim: q rows 512 (gated), k / v rows 128 +I, MI, E = 128, 128, 4 # shared / dense MLP width, routed expert width, routed experts +BLOCK = 128 +LM = "model.language_model" +FP8 = torch.float8_e4m3fn + +pytestmark = pytest.mark.filterwarnings("ignore::UserWarning") + + +@pytest.fixture(scope="session", autouse=True) +def _tp_info(): + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +def _bf16(*shape: int) -> torch.Tensor: + return torch.randn(*shape).to(torch.bfloat16) + + +# --------------------------------------------------------------------------- checkpoints + + +def _dense_bf16(moe: bool) -> dict[str, torch.Tensor]: + """Every non-expert tensor in bf16: layer 0 = GDN, layer 1 = attention, plus mtp / visual noise.""" + raw = {f"{LM}.embed_tokens.weight": _bf16(V, H), f"{LM}.norm.weight": _bf16(H), "lm_head.weight": _bf16(V, H)} + for layer in (0, 1): + pre = f"{LM}.layers.{layer}" + raw.update({f"{pre}.input_layernorm.weight": _bf16(H), f"{pre}.post_attention_layernorm.weight": _bf16(H)}) + mlp = f"{pre}.mlp.shared_expert" if moe else f"{pre}.mlp" + raw.update({f"{mlp}.gate_proj.weight": _bf16(I, H), f"{mlp}.up_proj.weight": _bf16(I, H), f"{mlp}.down_proj.weight": _bf16(H, I)}) + if moe: + raw.update({f"{pre}.mlp.gate.weight": _bf16(E, H), f"{pre}.mlp.shared_expert_gate.weight": _bf16(1, H)}) + gdn = f"{LM}.layers.0.linear_attn" + raw.update({ + f"{gdn}.in_proj_qkv.weight": _bf16(2 * KH * HD + VH * HD, H), f"{gdn}.in_proj_z.weight": _bf16(VH * HD, H), + f"{gdn}.in_proj_b.weight": _bf16(VH, H), f"{gdn}.in_proj_a.weight": _bf16(VH, H), + f"{gdn}.conv1d.weight": _bf16(2 * KH * HD + VH * HD, 1, 4), f"{gdn}.A_log": torch.randn(VH), + f"{gdn}.dt_bias": torch.randn(VH), f"{gdn}.norm.weight": _bf16(HD), f"{gdn}.out_proj.weight": _bf16(H, VH * HD), + }) + attn = f"{LM}.layers.1.self_attn" + raw.update({ + f"{attn}.q_proj.weight": _bf16(2 * QH * AHD, H), f"{attn}.k_proj.weight": _bf16(KVH * AHD, H), + f"{attn}.v_proj.weight": _bf16(KVH * AHD, H), f"{attn}.o_proj.weight": _bf16(H, QH * AHD), + f"{attn}.q_norm.weight": _bf16(AHD), f"{attn}.k_norm.weight": _bf16(AHD), + }) + raw.update({ + "mtp.layers.0.self_attn.q_proj.weight": _bf16(2 * QH * AHD, H), + "model.visual.blocks.0.attn.qkv.weight": _bf16(3 * H, H), + }) + return raw + + +def _nvfp4(weight: torch.Tensor, *, ct: bool) -> dict[str, torch.Tensor]: + """Packed NVFP4 tensors of one module under the ModelOpt or the llm-compressor names.""" + rows, cols = weight.shape + packed = torch.randint(0, 256, (rows, cols // 2), dtype=torch.uint8) + scale = (torch.rand(rows, cols // 16) + 0.5).to(FP8) + glob = torch.rand(1) + 0.5 + if ct: + return {"weight_packed": packed, "weight_scale": scale, "weight_global_scale": glob, "input_global_scale": torch.rand(1)} + return {"weight": packed, "weight_scale": scale, "weight_scale_2": glob.reshape(()), "input_scale": torch.rand(())} + + +def _fp8_tensor(weight: torch.Tensor) -> dict[str, torch.Tensor]: + """ModelOpt FP8: one fp32 scale for the whole weight plus the calibrated activation scale.""" + return {"weight": weight.to(FP8), "weight_scale": torch.rand(()) + 0.5, "input_scale": torch.rand(()) + 0.5} + + +def _fp8_tensor_ct(weight: torch.Tensor) -> dict[str, torch.Tensor]: + """llm-compressor ``strategy: tensor`` with static activations: bf16 scalar weight and input scales.""" + return {"weight": weight.to(FP8), "weight_scale": (torch.rand(1) + 0.5).to(torch.bfloat16), "input_scale": (torch.rand(1) + 0.5).to(torch.bfloat16)} + + +def _fp8_channel(weight: torch.Tensor) -> dict[str, torch.Tensor]: + """llm-compressor ``strategy: channel``: one bf16 scale per output row.""" + return {"weight": weight.to(FP8), "weight_scale": (torch.rand(weight.shape[0], 1) + 0.5).to(torch.bfloat16)} + + +def _fp8_block(weight: torch.Tensor, *, ct: bool) -> dict[str, torch.Tensor]: + """128x128 block-fp8: ``weight_scale_inv`` bf16 (HF fp8) or ``weight_scale`` fp32 (llm-compressor).""" + scale = torch.rand(weight.shape[0] // BLOCK, weight.shape[1] // BLOCK) + 0.5 + return {"weight": weight.to(FP8), "weight_scale": scale} if ct else {"weight": weight.to(FP8), "weight_scale_inv": scale.to(torch.bfloat16)} + + +def _quantize(raw: dict[str, torch.Tensor], modules: list[str], quantize) -> None: + for module in modules: + weight = raw.pop(f"{module}.weight") + raw.update({f"{module}.{suffix}": t for suffix, t in quantize(weight).items()}) + + +def _experts(raw: dict[str, torch.Tensor], quantize=None) -> None: + """Routed experts: stacked bf16 per layer, or per-expert tensors under ``quantize``.""" + for layer in (0, 1): + pre = f"{LM}.layers.{layer}.mlp.experts" + if quantize is None: + raw[f"{pre}.gate_up_proj"] = _bf16(E, 2 * MI, H) + raw[f"{pre}.down_proj"] = _bf16(E, H, MI) + continue + for expert in range(E): + for proj, shape in (("gate_proj", (MI, H)), ("up_proj", (MI, H)), ("down_proj", (H, MI))): + raw.update({f"{pre}.{expert}.{proj}.{suffix}": t for suffix, t in quantize(_bf16(*shape)).items()}) + + +GDN_QKVZ_OUT = [f"{LM}.layers.0.linear_attn.{p}" for p in ("in_proj_qkv", "in_proj_z", "out_proj")] +GDN_BA = [f"{LM}.layers.0.linear_attn.{p}" for p in ("in_proj_b", "in_proj_a")] +ATTN = [f"{LM}.layers.1.self_attn.{p}_proj" for p in "qkvo"] +SHARED = [f"{LM}.layers.{l}.mlp.shared_expert.{p}_proj" for l in (0, 1) for p in ("gate", "up", "down")] +DENSE_MLP = [f"{LM}.layers.{l}.mlp.{p}_proj" for l in (0, 1) for p in ("gate", "up", "down")] +ROUTERS = [f"{LM}.layers.{l}.mlp.{p}" for l in (0, 1) for p in ("gate", "shared_expert_gate")] + +NVFP4_GROUP = { + "weights": {"num_bits": 4, "type": "float", "strategy": "tensor_group", "group_size": 16, "symmetric": True}, + "input_activations": {"num_bits": 4, "type": "float", "strategy": "tensor_group", "group_size": 16, "dynamic": "local"}, + "format": "nvfp4-pack-quantized", +} +FP8_CHANNEL_GROUP = { + "weights": {"num_bits": 8, "type": "float", "strategy": "channel", "group_size": None, "symmetric": True}, + "input_activations": {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True}, + "format": "float-quantized", +} +FP8_TENSOR_STATIC_GROUP = { + "weights": {"num_bits": 8, "type": "float", "strategy": "tensor", "group_size": None, "symmetric": True}, + "input_activations": {"num_bits": 8, "type": "float", "strategy": "tensor", "dynamic": False}, + "format": "float-quantized", +} +FP8_BLOCK_GROUP = { + "weights": {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128], "symmetric": True}, + "input_activations": {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True}, + "format": "float-quantized", +} + + +def _ct(config_groups: dict, ignore: list[str], fmt: str) -> dict: + return {"quant_method": "compressed-tensors", "format": fmt, "config_groups": config_groups, "ignore": ignore, + "quantization_status": "compressed"} + + +# Qwen/Qwen3.6-35B-A3B-FP8: 128x128 block-fp8 everywhere but the listed modules +QWEN_FP8 = { + "quant_method": "fp8", "activation_scheme": "dynamic", "fmt": "e4m3", "weight_block_size": [128, 128], + "modules_to_not_convert": ["lm_head", "model.embed_tokens", *GDN_BA, *ROUTERS, + *[f"{LM}.layers.{l}.{n}" for l in (0, 1) for n in ("input_layernorm", "post_attention_layernorm")]], +} +# nvidia/Qwen3.6-35B-A3B-NVFP4: per-tensor FP8 attention / GDN, NVFP4 shared expert, experts and lm_head +MODELOPT_MIXED = { + "quant_method": "modelopt", "quant_algo": "MIXED_PRECISION", "ignore": ["mtp*"], + "quantized_layers": { + **{m: {"quant_algo": "FP8"} for m in GDN_QKVZ_OUT + ATTN}, + **{m: {"quant_algo": "W4A16_NVFP4", "group_size": 16} for m in SHARED + ["lm_head"]}, + **{f"{LM}.layers.{l}.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16} for l in (0, 1)}, + }, +} +# sakamakismile/Qwen3.6-27B-NVFP4: every Linear of the dense model, GDN in_proj included +CT_NVFP4_DENSE = _ct({"group_0": {**NVFP4_GROUP, "targets": ["Linear"]}}, ["lm_head"], "nvfp4-pack-quantized") +# RedHatAI/Qwen3.6-35B-A3B-NVFP4: every Linear but the GDN, the routers and lm_head +CT_NVFP4_MOE = _ct({"group_0": {**NVFP4_GROUP, "targets": ["Linear"]}}, ["lm_head", f"{LM}.embed_tokens", *GDN_QKVZ_OUT, *GDN_BA, *ROUTERS], "nvfp4-pack-quantized") +# unsloth/Qwen3.6-35B-A3B-NVFP4-Fast: channel-fp8 attention / GDN / lm_head, NVFP4 experts and shared expert; the ignore list names the linear_attn container too +CT_MIXED_FAST = _ct({ + "group_0": {**FP8_CHANNEL_GROUP, "targets": [r"re:.*self_attn\.(q|k|v|o)_proj$", r"re:.*linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$", "re:.*lm_head"]}, + "group_1": {**NVFP4_GROUP, "targets": [r"re:.*mlp\.experts\.\d+\.(gate|up|down)_proj$", r"re:.*shared_expert\.(gate|up|down)_proj$"]}, +}, [f"{LM}.layers.0.linear_attn", f"{LM}.layers.0.linear_attn.norm", *GDN_BA, *ROUTERS], "mixed-precision") +# primitive-ai/Ornith-1.5-35B-A3B-mixed-NVFP4-FP8: static per-tensor fp8 attention / GDN / shared expert, NVFP4 experts; the ignore list names every ``experts.N`` container +CT_TENSOR_FP8_MOE = _ct({ + "group_0": {**FP8_TENSOR_STATIC_GROUP, "targets": [r"re:.*\.self_attn\.(q_proj|k_proj|v_proj|o_proj)$", r"re:.*\.linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$", r"re:.*\.mlp\.shared_expert\.(gate_proj|up_proj|down_proj)$"]}, + "group_1": {**NVFP4_GROUP, "targets": [r"re:.*\.mlp\.experts\.\d+\.(gate_proj|up_proj|down_proj)$"]}, +}, ["lm_head", f"{LM}.layers.0.linear_attn", *GDN_BA, *ROUTERS, *[f"{LM}.layers.{l}.mlp.experts.{e}" for l in (0, 1) for e in range(E)]], "mixed-precision") +# block-fp8 dense side with NVFP4 experts (the JIAQI13 / kyaky export) +CT_BLOCK_MOE = _ct({ + "group_0": {**FP8_BLOCK_GROUP, "targets": [r"re:.*self_attn\.(q|k|v|o)_proj$", r"re:.*linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$", r"re:.*shared_expert\.(gate|up|down)_proj$"]}, + "group_1": {**NVFP4_GROUP, "targets": [r"re:.*mlp\.experts\.\d+\.(gate|up|down)_proj$"]}, +}, ["lm_head", *GDN_BA, *ROUTERS], "mixed-precision") + + +def _layout(name: str) -> tuple[bool, dict | None, dict[str, torch.Tensor]]: + """``(moe, quantization_config, raw tensors)`` of one released layout.""" + moe = name != "ct_nvfp4_dense" + raw = _dense_bf16(moe) + if name == "bf16": + _experts(raw) + return moe, None, raw + if name == "fp8_block": + _quantize(raw, GDN_QKVZ_OUT + ATTN + SHARED, lambda w: _fp8_block(w, ct=False)) + _experts(raw, lambda w: _fp8_block(w, ct=False)) + return moe, QWEN_FP8, raw + if name == "modelopt_mixed": + _quantize(raw, GDN_QKVZ_OUT + ATTN, _fp8_tensor) + _quantize(raw, SHARED + ["lm_head"], lambda w: _nvfp4(w, ct=False)) + _experts(raw, lambda w: _nvfp4(w, ct=False)) + return moe, MODELOPT_MIXED, raw + if name == "ct_nvfp4_dense": + _quantize(raw, GDN_QKVZ_OUT + GDN_BA + ATTN + DENSE_MLP, lambda w: _nvfp4(w, ct=True)) + return moe, CT_NVFP4_DENSE, raw + if name == "ct_nvfp4_moe": + _quantize(raw, ATTN + SHARED, lambda w: _nvfp4(w, ct=True)) + _experts(raw, lambda w: _nvfp4(w, ct=True)) + return moe, CT_NVFP4_MOE, raw + if name == "ct_mixed_fast": + _quantize(raw, GDN_QKVZ_OUT + ATTN + ["lm_head"], _fp8_channel) + _quantize(raw, SHARED, lambda w: _nvfp4(w, ct=True)) + _experts(raw, lambda w: _nvfp4(w, ct=True)) + raw.update({f"{LM}.layers.1.self_attn.k_scale": torch.rand(()), f"{LM}.layers.1.self_attn.v_scale": torch.rand(())}) + return moe, CT_MIXED_FAST, raw + if name == "ct_tensor_fp8_moe": + _quantize(raw, GDN_QKVZ_OUT + ATTN + SHARED, _fp8_tensor_ct) + _experts(raw, lambda w: _nvfp4(w, ct=True)) + return moe, CT_TENSOR_FP8_MOE, raw + if name == "ct_block_moe": + _quantize(raw, GDN_QKVZ_OUT + ATTN + SHARED, lambda w: _fp8_block(w, ct=True)) + _experts(raw, lambda w: _nvfp4(w, ct=True)) + return moe, CT_BLOCK_MOE, raw + raise KeyError(name) + + +LAYOUTS = ["bf16", "fp8_block", "modelopt_mixed", "ct_nvfp4_dense", "ct_nvfp4_moe", "ct_mixed_fast", "ct_tensor_fp8_moe", "ct_block_moe"] + + +def _config_json(moe: bool, quantization_config) -> dict: + text = { + "model_type": "qwen3_5_moe_text" if moe else "qwen3_5_text", "num_hidden_layers": 2, "hidden_size": H, "vocab_size": V, + "head_dim": AHD, "num_attention_heads": QH, "num_key_value_heads": KVH, "intermediate_size": I, + "layer_types": ["linear_attention", "full_attention"], + "rope_parameters": {"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 0.25}, + "max_position_embeddings": 4096, "rms_norm_eps": 1e-6, "hidden_act": "silu", "tie_word_embeddings": False, + "linear_num_key_heads": KH, "linear_num_value_heads": VH, "linear_key_head_dim": HD, "linear_value_head_dim": HD, + "linear_conv_kernel_dim": 4, + } + if moe: + text.update(num_experts=E, num_experts_per_tok=2, moe_intermediate_size=MI, shared_expert_intermediate_size=I) + return { + "architectures": ["Qwen3_5MoeForConditionalGeneration" if moe else "Qwen3_5ForConditionalGeneration"], + "model_type": "qwen3_5_moe" if moe else "qwen3_5", "text_config": text, "quantization_config": quantization_config, + } + + +def _write(folder, moe: bool, quantization_config, raw: dict[str, torch.Tensor], *, shards: int = 2) -> str: + """Spread the tensors over ``shards`` files, without an index, so fusions cross a file boundary.""" + names = sorted(raw) + for i in range(shards): + save_file({n: raw[n] for n in names[i::shards]}, str(folder / f"model-{i:05d}.safetensors")) + (folder / "config.json").write_text(json.dumps(_config_json(moe, quantization_config))) + return str(folder) + + +def _install(folder: str) -> None: + """Install the folder's QuantConfig process-wide, as EngineConfig does before the reader runs.""" + hf = cached_load_hf_config(folder) + set_quant_config(checkpoint_quant_config(folder, hf, get_model_spec(hf.architectures[0]))) + + +def _load(folder: str, *, experts: bool = False) -> dict[str, torch.Tensor]: + _install(folder) + return {n: t.clone() for n, t in iter_weights(folder, torch.device("cpu"), include_moe_experts=experts, include_non_moe=True)} + + +def _meta_state_dict(folder: str) -> dict[str, torch.Tensor]: + """State dict of the model the engine builds for ``folder`` (routed experts offloaded), on the meta device.""" + from freetoken.engine.config import EngineConfig + from freetoken.engine.engine import _decode_target + from freetoken.layers import rotary + from freetoken.models import create_model + from freetoken.utils.torch_utils import torch_dtype + + strategy = "offload" if cached_load_hf_config(folder).architectures[0].startswith("Qwen3_5Moe") else "auto" + config = EngineConfig(model_path=folder, tp_info=try_get_tp_info(), dtype=torch.bfloat16, moe_strategy=strategy) + object.__setattr__(config.model_config, "moe_strategy", strategy) + object.__setattr__(config.model_config, "decode_target", _decode_target(config)) + saved = rotary._ROPE_DEVICE + rotary.set_rope_device(torch.device("cpu")) # get_rope refuses to build on meta + rotary.get_rope.cache_clear() + try: + with torch.device("meta"), torch_dtype(torch.bfloat16): + return create_model(config.model_config).state_dict() + finally: + rotary.set_rope_device(saved) + rotary.get_rope.cache_clear() + + +@pytest.fixture(scope="module", params=LAYOUTS) +def checkpoint(request, tmp_path_factory): + torch.manual_seed(LAYOUTS.index(request.param)) + moe, quant, raw = _layout(request.param) + return request.param, _write(tmp_path_factory.mktemp(request.param), moe, quant, raw), raw + + +# --------------------------------------------------------------------------- the reader against the engine's model + + +def test_emitted_keys_are_the_model_state_dict(checkpoint): + """Every layout fills exactly the buffers the engine builds from the same config, with the buffers' shapes and (for the weights) dtypes.""" + _name, folder, _raw = checkpoint + loaded, state = _load(folder), _meta_state_dict(folder) + assert set(loaded) == set(state) + for key, tensor in loaded.items(): + assert tensor.shape == state[key].shape, key + if key.endswith(".weight"): + assert tensor.dtype is state[key].dtype, key + assert not any(k.endswith((".input_global_scale", ".weight_scale_2", ".weight_global_scale", ".k_scale")) for k in loaded) + assert not any(".mlp.experts." in k or k.startswith(("mtp.", "model.visual.")) for k in loaded) + + +def test_expert_quant_tag_follows_the_config(checkpoint): + name, folder, _raw = checkpoint + config = parse_config(cached_load_hf_config(folder)) + expected = {"bf16": "none", "fp8_block": "fp8_block", "ct_nvfp4_dense": "none"}.get(name, "nvfp4") + assert config.expert_quant == expected + assert config.weight_block_size == ((128, 128) if name == "fp8_block" else None) + + +def _slices(fused: torch.Tensor, parts: list[torch.Tensor]) -> list[torch.Tensor]: + return list(torch.split(fused, [p.shape[0] for p in parts], dim=0)) + + +def _same(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.dtype is b.dtype and torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + + +def test_bf16_fusions_slice_back_and_norms_get_plus_one(checkpoint): + name, folder, raw = checkpoint + if name != "bf16": + pytest.skip("bf16 layout only") + loaded = _load(folder) + gdn, attn = f"{LM}.layers.0.linear_attn", f"{LM}.layers.1.self_attn" + parts = [raw[f"{gdn}.in_proj_{p}.weight"] for p in ("qkv", "z", "b", "a")] + assert all(_same(p, s) for p, s in zip(parts, _slices(loaded["model.layers.0.linear_attn.in_proj.weight"], parts))) + parts = [raw[f"{attn}.{p}_proj.weight"] for p in "qkv"] + assert all(_same(p, s) for p, s in zip(parts, _slices(loaded["model.layers.1.self_attn.qkv_proj.weight"], parts))) + merged = loaded["model.layers.1.mlp.shared_expert.gate_up_proj.weight"] + assert _same(merged[:I], raw[f"{LM}.layers.1.mlp.shared_expert.gate_proj.weight"]) + assert _same(merged[I:], raw[f"{LM}.layers.1.mlp.shared_expert.up_proj.weight"]) + assert torch.equal(loaded["model.norm.weight"], raw[f"{LM}.norm.weight"] + 1.0) + assert torch.equal(loaded["model.layers.1.self_attn.q_norm.weight"], raw[f"{attn}.q_norm.weight"] + 1.0) + assert torch.equal(loaded["model.layers.0.linear_attn.norm.weight"], raw[f"{gdn}.norm.weight"]) + assert loaded["model.layers.0.linear_attn.A_log"].dtype is torch.float32 + + +def test_bf16_stacked_experts_pass_through_only_when_asked(checkpoint): + name, folder, raw = checkpoint + if name != "bf16": + pytest.skip("bf16 layout only") + _install(folder) + experts = dict(iter_weights(folder, torch.device("cpu"), include_moe_experts=True, include_non_moe=False)) + assert set(experts) == {f"model.layers.{l}.mlp.experts.{p}" for l in (0, 1) for p in ("gate_up_proj", "down_proj")} + assert torch.equal(experts["model.layers.0.mlp.experts.gate_up_proj"], raw[f"{LM}.layers.0.mlp.experts.gate_up_proj"]) + assert not any(".experts." in k for k in _load(folder)) + + +def test_block_fp8_fuses_weight_and_scale_per_kind(checkpoint): + name, folder, raw = checkpoint + if name not in ("fp8_block", "ct_block_moe"): + pytest.skip("block-fp8 layouts only") + loaded = _load(folder) + scale = "weight_scale_inv" if name == "fp8_block" else "weight_scale" + gdn, attn = f"{LM}.layers.0.linear_attn", f"{LM}.layers.1.self_attn" + for fused, sources in ( + ("model.layers.1.self_attn.qkv_proj", [f"{attn}.{p}_proj" for p in "qkv"]), + ("model.layers.0.linear_attn.in_proj_qkvz", [f"{gdn}.in_proj_qkv", f"{gdn}.in_proj_z"]), + ("model.layers.0.mlp.shared_expert.gate_up_proj", [f"{LM}.layers.0.mlp.shared_expert.{p}_proj" for p in ("gate", "up")]), + ): + parts = [raw[f"{s}.weight"] for s in sources] + assert all(_same(p, b) for p, b in zip(parts, _slices(loaded[f"{fused}.weight"], parts))) + scales = [raw[f"{s}.{scale}"] for s in sources] + assert all(_same(p, b) for p, b in zip(scales, _slices(loaded[f"{fused}.weight_scale_inv"], scales))) + assert loaded["model.layers.0.linear_attn.in_proj_qkvz.weight_scale_inv"].shape == (6, 1) + ba = loaded["model.layers.0.linear_attn.in_proj_ba.weight"] + assert ba.dtype is torch.bfloat16 and torch.equal(ba, torch.cat([raw[f"{gdn}.in_proj_b.weight"], raw[f"{gdn}.in_proj_a.weight"]])) + assert loaded["lm_head.weight"].dtype is torch.bfloat16 + + +def test_modelopt_fp8_scales_broadcast_per_part_and_input_scale_is_the_max(checkpoint): + name, folder, raw = checkpoint + if name != "modelopt_mixed": + pytest.skip("modelopt layout only") + loaded = _load(folder) + attn = f"{LM}.layers.1.self_attn" + scale = loaded["model.layers.1.self_attn.qkv_proj.weight_scale"] + assert scale.dtype is torch.float32 and scale.shape == (2 * QH * AHD + 2 * KVH * AHD,) + expected = torch.cat([raw[f"{attn}.{p}_proj.weight_scale"].expand(raw[f"{attn}.{p}_proj.weight"].shape[0]) for p in "qkv"]) + assert torch.equal(scale, expected) + assert torch.equal(loaded["model.layers.1.self_attn.qkv_proj.input_scale"], torch.stack([raw[f"{attn}.{p}_proj.input_scale"] for p in "qkv"]).max()) + assert loaded["model.layers.1.self_attn.o_proj.input_scale"].shape == () + # NVFP4 shared expert: each part keeps its own block scales and global; the input scale is the max of the parts + shared = f"{LM}.layers.0.mlp.shared_expert" + fused = "model.layers.0.mlp.shared_expert.gate_up_proj" + assert _same(loaded[f"{fused}.weight"][:I], raw[f"{shared}.gate_proj.weight"]) + assert _same(loaded[f"{fused}.weight_scale"][I:], raw[f"{shared}.up_proj.weight_scale"]) + glob = loaded[f"{fused}.weight_global"] + assert glob.dtype is torch.float16 and torch.equal(glob[:I], raw[f"{shared}.gate_proj.weight_scale_2"].to(torch.float16).expand(I)) + assert torch.equal(loaded[f"{fused}.input_scale"], torch.stack([raw[f"{shared}.{p}_proj.input_scale"] for p in ("gate", "up")]).max()) + assert loaded["lm_head.input_scale"].shape == () + assert loaded["lm_head.weight"].dtype is torch.uint8 and loaded["lm_head.weight_global"].shape == (V,) + + +def test_compressed_tensors_nvfp4_global_is_the_reciprocal(checkpoint): + name, folder, raw = checkpoint + if name not in ("ct_nvfp4_dense", "ct_nvfp4_moe"): + pytest.skip("llm-compressor NVFP4 layouts only") + loaded = _load(folder) + attn = f"{LM}.layers.1.self_attn" + glob = loaded["model.layers.1.self_attn.qkv_proj.weight_global"] + expected = torch.cat([(1.0 / raw[f"{attn}.{p}_proj.weight_global_scale"]).to(torch.float16).expand(raw[f"{attn}.{p}_proj.weight_packed"].shape[0]) for p in "qkv"]) + assert glob.dtype is torch.float16 and torch.equal(glob, expected) + assert _same(loaded["model.layers.1.self_attn.o_proj.weight_scale"], raw[f"{attn}.o_proj.weight_scale"]) + # ``dynamic: local`` ships the quant-side activation global too; it lands as the dequant-side input_scale + assert torch.equal(loaded["model.layers.1.self_attn.o_proj.input_scale"], (1.0 / raw[f"{attn}.o_proj.input_global_scale"]).reshape(())) + assert torch.equal(loaded["model.layers.1.self_attn.qkv_proj.input_scale"], torch.stack([1.0 / raw[f"{attn}.{p}_proj.input_global_scale"].reshape(()) for p in "qkv"]).max()) + assert loaded["lm_head.weight"].dtype is torch.bfloat16 + gdn = f"{LM}.layers.0.linear_attn" + if name == "ct_nvfp4_dense": + # the GDN projections are quantized too, so qkv|z and b|a are both native NVFP4 + assert loaded["model.layers.0.linear_attn.in_proj_qkvz.weight"].dtype is torch.uint8 + ba = loaded["model.layers.0.linear_attn.in_proj_ba.weight"] + assert ba.dtype is torch.uint8 and _same(ba[VH:], raw[f"{gdn}.in_proj_a.weight_packed"]) + assert loaded["model.layers.0.linear_attn.in_proj_ba.weight_global"].shape == (2 * VH,) + else: + parts = [raw[f"{gdn}.in_proj_{p}.weight"] for p in ("qkv", "z", "b", "a")] + assert all(_same(p, s) for p, s in zip(parts, _slices(loaded["model.layers.0.linear_attn.in_proj.weight"], parts))) + + +def test_channel_fp8_keeps_row_order_and_an_ignored_container_does_not_shield_its_children(checkpoint): + name, folder, raw = checkpoint + if name != "ct_mixed_fast": + pytest.skip("unsloth layout only") + loaded = _load(folder) + attn, gdn = f"{LM}.layers.1.self_attn", f"{LM}.layers.0.linear_attn" + scale = loaded["model.layers.1.self_attn.qkv_proj.weight_scale"] + assert scale.dtype is torch.float32 + assert torch.equal(scale, torch.cat([raw[f"{attn}.{p}_proj.weight_scale"].reshape(-1) for p in "qkv"]).to(torch.float32)) + assert "model.layers.1.self_attn.qkv_proj.input_scale" not in loaded # dynamic per-token fp8: no static activation scale + # ``ignore`` names ``layers.0.linear_attn`` itself; its quantized projections still load as fp8 + qkvz = loaded["model.layers.0.linear_attn.in_proj_qkvz.weight"] + assert qkvz.dtype is FP8 and _same(qkvz[: 2 * KH * HD + VH * HD], raw[f"{gdn}.in_proj_qkv.weight"]) + assert loaded["model.layers.0.linear_attn.in_proj_ba.weight"].dtype is torch.bfloat16 + assert loaded["lm_head.weight"].dtype is FP8 and loaded["lm_head.weight_scale"].shape == (V,) + assert loaded["model.layers.0.mlp.shared_expert.down_proj.weight"].dtype is torch.uint8 + + +def test_static_fp8_keeps_the_activation_scale_and_ignored_expert_containers_do_not_shield_the_experts(checkpoint): + name, folder, raw = checkpoint + if name != "ct_tensor_fp8_moe": + pytest.skip("Ornith layout only") + loaded = _load(folder) + attn = f"{LM}.layers.1.self_attn" + assert loaded["model.layers.1.self_attn.qkv_proj.weight_scale"].shape == (2 * QH * AHD + 2 * KVH * AHD,) + expected = torch.stack([raw[f"{attn}.{p}_proj.input_scale"].reshape(()) for p in "qkv"]).max().to(torch.float32) + assert torch.equal(loaded["model.layers.1.self_attn.qkv_proj.input_scale"], expected) + assert loaded["model.layers.0.mlp.shared_expert.down_proj.input_scale"].dtype is torch.float32 + assert loaded["lm_head.weight"].dtype is torch.bfloat16 + assert parse_config(cached_load_hf_config(folder)).expert_quant == "nvfp4" + + +# --------------------------------------------------------------------------- the expert reader shares the dialect names + + +def test_nvfp4_expert_pieces_read_either_dialect_from_a_single_file(checkpoint): + name, folder, raw = checkpoint + if name not in ("modelopt_mixed", "ct_nvfp4_moe", "ct_tensor_fp8_moe"): + pytest.skip("NVFP4 expert layouts only") + _install(folder) + config = parse_config(cached_load_hf_config(folder)) + spec = nvfp4_expert_spec(folder, config) + pieces = list(iter_nvfp4_expert_pieces(folder, config, spec)) + assert len(pieces) == 2 * E + layer, e0, e1, piece = next(p for p in pieces if p[0] == 1 and p[1] == 2) + assert (e0, e1) == (2, 3) + base = f"{LM}.layers.1.mlp.experts.2.gate_proj" + if name == "modelopt_mixed": + assert _same(piece["gate"][0], raw[f"{base}.weight"]) + assert torch.equal(piece["gate_global"].reshape(-1), raw[f"{base}.weight_scale_2"].reshape(-1).to(torch.float16)) + else: + assert _same(piece["gate"][0], raw[f"{base}.weight_packed"]) + assert _same(piece["down_scale"][0], raw[f"{LM}.layers.1.mlp.experts.2.down_proj.weight_scale"]) + assert torch.equal(piece["gate_global"].reshape(-1), (1.0 / raw[f"{base}.weight_global_scale"]).to(torch.float16)) + + +# --------------------------------------------------------------------------- the family's unquantized modules + + +def _nvfp4_reference(packed: torch.Tensor, scale: torch.Tensor, global_rows: torch.Tensor) -> torch.Tensor: + """e2m1 codes (low nibble first) x e4m3 block scale x per-row global, as the kernel computes it.""" + table = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + codes = torch.stack([packed & 0xF, packed >> 4], dim=-1).reshape(packed.shape[0], -1).long() + values = table[codes & 7] * torch.where(codes & 8 > 0, -1.0, 1.0) + return (values * scale.float().repeat_interleave(16, dim=1) * global_rows.float()[:, None]).to(torch.bfloat16) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the NVFP4 dequant kernel runs on CUDA") +def test_a_quantized_router_is_dequantized_because_the_family_serves_it_bf16(tmp_path): + """``targets: ["Linear"]`` with no ignore entry for ``mlp.gate`` quantizes the router; the model builds it bf16 (register.py unquantized_modules), so the reader dequantizes.""" + torch.manual_seed(11) + moe, quant, raw = _layout("ct_nvfp4_moe") + quant = {**quant, "ignore": [i for i in quant["ignore"] if not i.endswith(".mlp.gate")]} + gates = [f"{LM}.layers.{l}.mlp.gate" for l in (0, 1)] + _quantize(raw, gates, lambda w: _nvfp4(w, ct=True)) + loaded = _load(_write(tmp_path, moe, quant, raw)) + for layer, gate in enumerate(gates): + got = loaded[f"model.layers.{layer}.mlp.gate.weight"] + assert got.dtype is torch.bfloat16 and got.shape == (E, H) + expected = _nvfp4_reference(raw[f"{gate}.weight_packed"], raw[f"{gate}.weight_scale"], (1.0 / raw[f"{gate}.weight_global_scale"]).to(torch.float16).expand(E)) + assert torch.equal(got, expected) + + +# --------------------------------------------------------------------------- checkpoints that disagree with their config + + +ATTN1 = f"{LM}.layers.1.self_attn" +REJECTED = [ + pytest.param(None, lambda w: {f"{ATTN1}.q_proj.weight": w, f"{ATTN1}.q_proj.weight_scale": torch.rand(())}, + r"q_proj\.weight_scale: .*declares .*q_proj unquantized", id="scale the config does not declare"), + pytest.param(None, lambda w: {f"{ATTN1}.o_proj.weight": w.to(FP8)}, + r"o_proj\.weight is torch\.float8", id="fp8 weight the config declares bf16"), + pytest.param(MODELOPT_MIXED, lambda w: {f"{ATTN1}.o_proj.weight": w, f"{ATTN1}.o_proj.weight_scale": torch.rand(()), f"{ATTN1}.o_proj.input_scale": torch.rand(())}, + r"o_proj: weight is torch\.bfloat16", id="bf16 weight the config declares fp8"), + pytest.param(CT_MIXED_FAST, lambda w: {f"{ATTN1}.o_proj.weight": w.to(FP8), f"{ATTN1}.o_proj.weight_scale": torch.rand(3, 1)}, + "expected 1 or", id="per-channel scale with the wrong row count"), + pytest.param(CT_BLOCK_MOE, lambda w: {f"{ATTN1}.o_proj.weight": w.to(FP8), f"{ATTN1}.o_proj.weight_scale": torch.rand(2, 1)}, + r"weight_scale_inv is \(2, 1\), expected \(1, 2\)", id="block scale of the wrong shape"), + pytest.param(CT_NVFP4_MOE, lambda w: {f"{ATTN1}.o_proj.weight_packed": torch.zeros(H, QH * AHD // 2, dtype=torch.uint8), f"{ATTN1}.o_proj.weight_global_scale": torch.rand(1)}, + "missing tensors of .*o_proj", id="quantized module without its block scale"), +] + + +@pytest.mark.parametrize("quantization_config, tensors, match", REJECTED) +def test_a_checkpoint_disagreeing_with_its_quant_config_is_rejected(tmp_path, quantization_config, tensors, match): + save_file(tensors(_bf16(H, QH * AHD)), str(tmp_path / "model.safetensors")) + (tmp_path / "config.json").write_text(json.dumps(_config_json(True, quantization_config))) + with pytest.raises(ValueError, match=match): + _load(str(tmp_path)) From b52216c3df070f40918963806cf9c1c07210dc48 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Fri, 11 Sep 2026 05:24:34 +0000 Subject: [PATCH 2/2] fix(quant): address scheme reader review: hotfix dialect names, gemma4 cross-shard scales, ct block-fp8 experts --- python/freetoken/models/gemma4/weight.py | 153 +++++++++--------- python/freetoken/models/qwen3_5_moe/weight.py | 14 +- scripts/ftw_hotfix.py | 65 +++++--- tests/models/test_qwen3_5_moe_weight.py | 36 ++++- 4 files changed, 155 insertions(+), 113 deletions(-) diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index 0d2056a63..740f4660f 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -9,6 +9,7 @@ from freetoken.models.config import FullAttentionGroupConfig from freetoken.models.loader import ( MergeRule, + ShardReader, drop_page_cache, iter_weight_files, ) @@ -64,26 +65,26 @@ _NVFP4_DENSE_MLP_RE = re.compile(r"\.mlp\.(gate_proj|up_proj|down_proj)\.weight$") -def _nvfp4_dense_parts(f, raw_base: str, keyset: set[str]): +def _nvfp4_dense_parts(reader: ShardReader, raw_base: str): """Load an NVFP4 dense weight as the NVFP4 linear method's buffers: weight uint8 [O, IN//2], weight_scale fp8-e4m3 block [O, IN//16], weight_global fp16 [O] (the per-tensor weight_scale_2 per output row), input_scale fp32 scalar or None when the export has none.""" - w = f.get_tensor(raw_base + ".weight") - s = f.get_tensor(raw_base + ".weight_scale") - g = f.get_tensor(raw_base + ".weight_scale_2").reshape(1).to(torch.float16) + w = reader.get_tensor(raw_base + ".weight") + s = reader.get_tensor(raw_base + ".weight_scale") + g = reader.get_tensor(raw_base + ".weight_scale_2").reshape(1).to(torch.float16) g = g.expand(w.shape[0]).contiguous() assert ( w.dtype is torch.uint8 and s.dtype is torch.float8_e4m3fn and g.dtype is torch.float16 ), f"unexpected NVFP4 dense dtypes at {raw_base}: {w.dtype}/{s.dtype}/{g.dtype}" - a = f.get_tensor(raw_base + ".input_scale").reshape(()).to(torch.float32) if raw_base + ".input_scale" in keyset else None + a = reader.get_tensor(raw_base + ".input_scale").reshape(()).to(torch.float32) if reader.has(raw_base + ".input_scale") else None return w, s, g, a -def _emit_nvfp4_dense_mlp(f, base: str, raw_base: str, buf: dict, keyset: set[str]): +def _emit_nvfp4_dense_mlp(reader: ShardReader, base: str, raw_base: str, buf: dict): """(key, tensor) pairs for an NVFP4 dense MLP projection: down_proj standalone; gate_proj/up_proj merged output-wise into gate_up_proj (each keeps its own scales, so the fused weight is exact). Returns [] while a gate/up merge is still buffered.""" - w, s, g, a = _nvfp4_dense_parts(f, raw_base, keyset) + w, s, g, a = _nvfp4_dense_parts(reader, raw_base) if base.endswith(".down_proj"): out = [(base + ".weight", w), (base + ".weight_scale", s), (base + ".weight_global", g)] return out + ([(base + ".input_scale", a)] if a is not None else []) @@ -168,77 +169,79 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None: } merge_buf: dict[str, dict[str, torch.Tensor]] = {} gateup_buf: dict[str, dict[str, tuple]] = {} - for file in tqdm( - iter_weight_files(model_path), - desc="Loading weights", - disable=not tp_info.is_primary(), - ): - with safetensors.safe_open(file, framework="pt", device=str(device)) as f: - keyset = set(f.keys()) - for raw_name in f.keys(): - name = rename_key(raw_name, include_vision=include_vision) - if name is None: - continue - - # Per-expert NVFP4 tensors go to the offload cache (expert pieces), - # not this dense pass; fused bf16/q4_0 experts lack ".experts.." so are unaffected. - if _NVFP4_EXPERT_RE.search(raw_name): - continue - - # NVFP4 dense-MLP scales are consumed with their .weight (below), never yielded. - if raw_name.endswith(_NVFP4_DENSE_SCALE_SUFFIXES): - continue - - is_vision = name.startswith(("vision_tower.", "embed_vision.")) - is_expert = ( - not is_vision and _PACKED_EXPERT_PATTERN.match(name) is not None - ) - if is_expert and not include_moe_experts: - continue - if not is_expert and not include_non_moe: - continue - - # Native W4A16 NVFP4 dense MLP: the .weight is FP4-packed and carries block + - # per-tensor scales. The keyset guard (weight_scale_2 sibling present) is - # defense-in-depth beyond config.dense_quant -- the sibling MoE checkpoint's - # bf16 shared_mlp has no such sibling, so it falls through to the bf16 path. - if ( - config.dense_quant == "nvfp4" - and not is_vision - and not is_expert - and _NVFP4_DENSE_MLP_RE.search(raw_name) - and raw_name[: -len(".weight")] + ".weight_scale_2" in keyset - ): - yield from _emit_nvfp4_dense_mlp( - f, name[: -len(".weight")], raw_name[: -len(".weight")], gateup_buf, keyset + reader = ShardReader(model_path, device) + try: + for file in tqdm( + iter_weight_files(model_path), + desc="Loading weights", + disable=not tp_info.is_primary(), + ): + with safetensors.safe_open(file, framework="pt", device=str(device)) as f: + for raw_name in f.keys(): + name = rename_key(raw_name, include_vision=include_vision) + if name is None: + continue + + # Per-expert NVFP4 tensors go to the offload cache (expert pieces), + # not this dense pass; fused bf16/q4_0 experts lack ".experts.." so are unaffected. + if _NVFP4_EXPERT_RE.search(raw_name): + continue + + # NVFP4 dense-MLP scales are consumed with their .weight (below), never yielded. + if raw_name.endswith(_NVFP4_DENSE_SCALE_SUFFIXES): + continue + + is_vision = name.startswith(("vision_tower.", "embed_vision.")) + is_expert = ( + not is_vision and _PACKED_EXPERT_PATTERN.match(name) is not None ) - continue - - tensor = f.get_tensor(raw_name) - if is_vision or is_expert: - yield name, tensor - continue - - info = merge_info(name) - if info is None: - yield name, tensor - continue - - merged_key, rule = info - slots = merge_buf.setdefault(merged_key, {}) - slots[rule.slot] = tensor - if rule.slot == "k" and k_eq_v_layers: - layer_match = _LAYER_INDEX_PATTERN.search(name) + if is_expert and not include_moe_experts: + continue + if not is_expert and not include_non_moe: + continue + + # Native W4A16 NVFP4 dense MLP: the .weight is FP4-packed and carries block + per-tensor scales. + # The weight_scale_2 sibling guard is defense-in-depth beyond config.dense_quant -- the sibling MoE checkpoint's bf16 shared_mlp has no such sibling, so it falls through to the bf16 path. if ( - layer_match is not None - and int(layer_match.group(1)) in k_eq_v_layers + config.dense_quant == "nvfp4" + and not is_vision + and not is_expert + and _NVFP4_DENSE_MLP_RE.search(raw_name) + and reader.has(raw_name[: -len(".weight")] + ".weight_scale_2") ): - slots["v"] = tensor - if not all(slot in slots for slot in rule.slots): - continue - parts = [slots[slot] for slot in rule.slots] - del merge_buf[merged_key] - yield merged_key, torch.cat(parts, dim=0) + yield from _emit_nvfp4_dense_mlp( + reader, name[: -len(".weight")], raw_name[: -len(".weight")], gateup_buf + ) + continue + + tensor = f.get_tensor(raw_name) + if is_vision or is_expert: + yield name, tensor + continue + + info = merge_info(name) + if info is None: + yield name, tensor + continue + + merged_key, rule = info + slots = merge_buf.setdefault(merged_key, {}) + slots[rule.slot] = tensor + if rule.slot == "k" and k_eq_v_layers: + layer_match = _LAYER_INDEX_PATTERN.search(name) + if ( + layer_match is not None + and int(layer_match.group(1)) in k_eq_v_layers + ): + slots["v"] = tensor + if not all(slot in slots for slot in rule.slots): + continue + parts = [slots[slot] for slot in rule.slots] + del merge_buf[merged_key] + yield merged_key, torch.cat(parts, dim=0) + + finally: + reader.close() assert not merge_buf, f"Incomplete merge groups in checkpoint: {list(merge_buf.keys())}" assert not gateup_buf, f"Incomplete NVFP4 gate/up merges: {list(gateup_buf.keys())}" diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d73df9bab..5fc12e747 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -303,9 +303,9 @@ def _is_expert(raw_name: str) -> bool: # Routed-expert checkpoint key (per-expert, un-fused). ``mtp.layers...`` is excluded by the # ``model.language_model.`` anchor, so the parallel reader only sees the real experts. -_FP8_EXPERT_RE = re.compile( +_FP8_EXPERT_KEY_RE = ( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate|up|down)_proj\.(?Pweight|weight_scale_inv)$" + r"(?Pgate|up|down)_proj\.(?Pweight|{scale})$" ) @@ -348,7 +348,7 @@ def _moe_dims(model_config): def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): """Block-fp8 routed experts, one piece per expert: ``{gate, up, down}`` fp8 codes and their - ``_scale`` (bf16 block ``weight_scale_inv``) companions. Other expert kinds use the generic readers.""" + ``_scale`` (block scale) companions, named as the checkpoint's dialect stores them. Other expert kinds use the generic readers.""" if kind is not QuantKind.FP8_BLOCK: return None if get_tp_info().size > 1: @@ -357,10 +357,12 @@ def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | from freetoken.moe.expert_pieces import per_expert_pieces L, E, H, I, dense = _moe_dims(config) - suffix = {"weight": "", "weight_scale_inv": "_scale"} + scale = get_quant_config().stored_tensors(QuantKind.FP8_BLOCK)["weight_scale_inv"].name + key_re = re.compile(_FP8_EXPERT_KEY_RE.format(scale=re.escape(scale))) + suffix = {"weight": "", scale: "_scale"} def locate(raw_name: str): - m = _FP8_EXPERT_RE.match(raw_name) + m = key_re.match(raw_name) if m is None: return None li = int(m["layer"]) - dense @@ -372,7 +374,7 @@ def locate(raw_name: str): parallel = experts_scattered(model_path) if parallel: tensors = iter_expert_tensors_parallel( - model_path, lambda n: _FP8_EXPERT_RE.match(n) is not None, workers=workers, chunk=chunk + model_path, lambda n: key_re.match(n) is not None, workers=workers, chunk=chunk ) return per_expert_pieces(tensors, locate, tensors_per_expert=6) diff --git a/scripts/ftw_hotfix.py b/scripts/ftw_hotfix.py index d39ed8361..7e324ed8a 100644 --- a/scripts/ftw_hotfix.py +++ b/scripts/ftw_hotfix.py @@ -41,6 +41,7 @@ from freetoken.engine.config import EngineConfig from freetoken.engine.engine import _decode_target from freetoken.layers import set_rope_device +from freetoken.layers.quantization.configs.base import Stored from freetoken.layers.quantization.names import NameMap from freetoken.models import create_model from freetoken.models.register import get_model_spec @@ -171,9 +172,21 @@ def expected_tensors(ftw_dir: str, resident_experts: bool): model = create_model(cfg.model_config) state = {k: (tuple(v.shape), v.dtype) for k, v in model.state_dict().items()} arch = cfg.model_config.architectures[0] - spec = get_model_spec(arch) - name_map = NameMap(roots=spec.checkpoint_roots, segments=spec.checkpoint_segments, packed=spec.packed_modules_mapping) - return arch, state, name_map + quant = cfg.model_config.quant + if quant is None: + spec = get_model_spec(arch) + name_map = NameMap(roots=spec.checkpoint_roots, segments=spec.checkpoint_segments, packed=spec.packed_modules_mapping) + else: + name_map = quant.name_map + return arch, state, name_map, quant + + +def stored_entry(quant, module: str, role: str) -> Stored: + """The checkpoint tensor behind ``module``'s ``role`` in the checkpoint's dialect; a plain same-named tensor when the dialect has no say.""" + scheme = quant.scheme_for(module) if quant is not None else None + if scheme is None: + return Stored(role) + return quant.storage(scheme).get(role, Stored(role)) # ------------------------------------------------------------------ repairs @@ -185,7 +198,7 @@ def dsv4_rename(name: str) -> str: return "model." + name -def plan(arch: str, entries: list[dict], expected: dict, name_map: NameMap): +def plan(arch: str, entries: list[dict], expected: dict, name_map: NameMap, quant=None): """Return (renames, dequants, fetches, drops, leftovers) that turn the FTW dense set into ``expected``.""" dense = {e["name"]: e for e in entries if e["kind"] == "weight"} renames: dict[str, str] = {} @@ -207,29 +220,28 @@ def plan(arch: str, entries: list[dict], expected: dict, name_map: NameMap): dequants.append((n, e, scale)) drops.add(scale["name"]) - # (name, checkpoint parts): a fused module maps to several checkpoint tensors - fetches: list[tuple[str, list[str]]] = [] + # (name, checkpoint parts, reciprocal): a fused module maps to several checkpoint tensors, named as the dialect stores the role + fetches: list[tuple[str, list[str], bool]] = [] for n in (n for n in expected if n not in names): module, _, leaf = n.rpartition(".") - parts = [f"{m}.{leaf}" for m in name_map.to_checkpoint(module)] if module else [n] - fetches.append((n, parts)) + entry = stored_entry(quant, module, leaf) if module else Stored(leaf) + parts = [f"{m}.{entry.name}" for m in name_map.to_checkpoint(module)] if module else [n] + fetches.append((n, parts, entry.reciprocal)) leftovers = [n for n in names if n not in expected and n not in drops] return renames, dequants, fetches, drops, leftovers -def resolve_fetches(fetches, source: TensorSource) -> tuple[dict[str, list[str]], list[str]]: - """Map each missing tensor to the source tensors it is built from; a fused scale needs every part.""" - resolved: dict[str, list[str]] = {} +def resolve_fetches(fetches, source: TensorSource) -> tuple[dict[str, tuple[list[str], bool]], list[str]]: + """Map each missing tensor to (the source tensors it is built from, reciprocal); a fused scale needs every part.""" + resolved: dict[str, tuple[list[str], bool]] = {} errors: list[str] = [] - for n, parts in fetches: - if source.has(n): - resolved[n] = [n] - elif all(source.has(p) for p in parts): + for n, parts, reciprocal in fetches: + if all(source.has(p) for p in parts): if len(parts) > 1 and not n.endswith(".input_scale"): errors.append(f"{n} maps to {len(parts)} source tensors; fusing is not supported here") else: - resolved[n] = parts + resolved[n] = (parts, reciprocal) else: errors.append(f"no source tensor for {n}: missing {[p for p in parts if not source.has(p)]}") return resolved, errors @@ -748,10 +760,10 @@ def main(argv: list[str] | None = None) -> int: resident_experts = any(e["kind"] == "weight" and ".experts." in e["name"] for e in index["tensors"]) log(f"reading {os.path.join(ns.ftw, INDEX_NAME)}: {len(index['tensors'])} entries, {len(index['shards'])} shards, {index['total_bytes'] / 2**30:.2f} GiB") log("building the current model on the meta device from the FTW's config.json" + (" (resident experts)" if resident_experts else "")) - arch, expected, name_map = expected_tensors(ns.ftw, resident_experts) + arch, expected, name_map, quant = expected_tensors(ns.ftw, resident_experts) log(f"{arch}: model declares {len(expected)} dense tensors") source = TensorSource(ns.repo, ns.source, ns.revision) if (ns.repo or ns.source) else None - renames, dequants, fetches, drops, leftovers = plan(arch, index["tensors"], expected, name_map) + renames, dequants, fetches, drops, leftovers = plan(arch, index["tensors"], expected, name_map, quant) dequant_names = {n for n, _, _ in dequants} ple = PleSpec(ns.ftw) if arch.startswith("Qwen4Exp") else None chk = check_ftw(ns.ftw, index, expected, renames, dequant_names, ple) @@ -764,7 +776,7 @@ def main(argv: list[str] | None = None) -> int: print(f" renames {len(renames)} dequantize {len(dequants)} fetch {len(fetches)} drop {len(drops)} leftover {len(leftovers)}" + (f" PLE table: {ple_status}" + (f" ({ple_detail})" if ple_detail else "") if need_ple else "") + (f" dead bytes in {len(dirty)} shard(s)" if dirty else "")) - for n, parts in fetches[:8]: + for n, parts, _ in fetches[:8]: print(f" fetch {n} <- {parts}") if len(fetches) > 8: print(f" ... {len(fetches) - 8} more") @@ -790,7 +802,7 @@ def main(argv: list[str] | None = None) -> int: if (fetches or need_ple) and source is None: print("ERROR: tensors must be fetched but neither --repo nor --source was given", file=sys.stderr) return 2 - fetch_srcs: dict[str, list[str]] = {} + fetch_srcs: dict[str, tuple[list[str], bool]] = {} if fetches: fetch_srcs, errors = resolve_fetches(fetches, source) for msg in errors: @@ -798,7 +810,7 @@ def main(argv: list[str] | None = None) -> int: bad = True # the converter transforms most tensors on the way in (norm offsets, fusion, packing); only the # activation scale scalars are stored as the checkpoint has them, so only they can be fetched raw - unfetchable = [n for n, _ in fetches if not n.endswith(".input_scale")] + unfetchable = [n for n, *_ in fetches if not n.endswith(".input_scale")] if unfetchable: print(f"ERROR: {len(unfetchable)} missing tensor(s) cannot be fetched raw (first: {unfetchable[0]}); reconvert the checkpoint", file=sys.stderr) bad = True @@ -863,11 +875,14 @@ def write_new(w: ShardWriter) -> None: tick(b, e["nbytes"]) done(b) items = list(fetch_srcs.items()) - for n, srcs in (items if _VERBOSE or not items else count_bar(items, "Fetching tensors")): - log(f" fetch {n} <- {', '.join(srcs)}") - vals = [source.get(c) for c in srcs] + for n, (srcs, reciprocal) in (items if _VERBOSE or not items else count_bar(items, "Fetching tensors")): + log(f" fetch {n} <- {', '.join(srcs)}" + (" (reciprocal)" if reciprocal else "")) + vals = [source.get(c).reshape(()).float() for c in srcs] + # llm-compressor stores the quant-side global; the layer wants the dequant-side scale, as the reader loads it + if reciprocal: + vals = [1.0 / v for v in vals] # a fused projection shares one activation scale; the reader takes the max over its parts - w.add_tensor(n, torch.stack([v.reshape(()).float() for v in vals]).max().reshape(())) + w.add_tensor(n, torch.stack(vals).max().reshape(())) hotfix = {"from": os.path.abspath(ns.ftw), "renamed": len(renames), "dequantized": len(dequants), "fetched": len(fetch_srcs), "dropped": len(drops), "compacted": 0, "source": ns.repo or ns.source} diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py index 98f21649c..1efec609c 100644 --- a/tests/models/test_qwen3_5_moe_weight.py +++ b/tests/models/test_qwen3_5_moe_weight.py @@ -8,17 +8,16 @@ from __future__ import annotations import json -import re import pytest import torch from safetensors.torch import save_file from freetoken.distributed import set_tp_info, try_get_tp_info -from freetoken.layers.quantization import set_quant_config +from freetoken.layers.quantization import QuantKind, set_quant_config from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces from freetoken.models.qwen3_5_moe.config import parse_config -from freetoken.models.qwen3_5_moe.weight import iter_weights, nvfp4_expert_spec +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.utils import cached_load_hf_config @@ -195,6 +194,8 @@ def _ct(config_groups: dict, ignore: list[str], fmt: str) -> dict: "group_0": {**FP8_BLOCK_GROUP, "targets": [r"re:.*self_attn\.(q|k|v|o)_proj$", r"re:.*linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$", r"re:.*shared_expert\.(gate|up|down)_proj$"]}, "group_1": {**NVFP4_GROUP, "targets": [r"re:.*mlp\.experts\.\d+\.(gate|up|down)_proj$"]}, }, ["lm_head", *GDN_BA, *ROUTERS], "mixed-precision") +# block-fp8 everywhere, experts included, under llm-compressor's names (``weight_scale`` for the block scale) +CT_BLOCK_EXPERTS = _ct({"group_0": {**FP8_BLOCK_GROUP, "targets": ["Linear"]}}, ["lm_head", f"{LM}.embed_tokens", *GDN_BA, *ROUTERS], "float-quantized") def _layout(name: str) -> tuple[bool, dict | None, dict[str, torch.Tensor]]: @@ -234,10 +235,14 @@ def _layout(name: str) -> tuple[bool, dict | None, dict[str, torch.Tensor]]: _quantize(raw, GDN_QKVZ_OUT + ATTN + SHARED, lambda w: _fp8_block(w, ct=True)) _experts(raw, lambda w: _nvfp4(w, ct=True)) return moe, CT_BLOCK_MOE, raw + if name == "ct_block_experts": + _quantize(raw, GDN_QKVZ_OUT + ATTN + SHARED, lambda w: _fp8_block(w, ct=True)) + _experts(raw, lambda w: _fp8_block(w, ct=True)) + return moe, CT_BLOCK_EXPERTS, raw raise KeyError(name) -LAYOUTS = ["bf16", "fp8_block", "modelopt_mixed", "ct_nvfp4_dense", "ct_nvfp4_moe", "ct_mixed_fast", "ct_tensor_fp8_moe", "ct_block_moe"] +LAYOUTS = ["bf16", "fp8_block", "modelopt_mixed", "ct_nvfp4_dense", "ct_nvfp4_moe", "ct_mixed_fast", "ct_tensor_fp8_moe", "ct_block_moe", "ct_block_experts"] def _config_json(moe: bool, quantization_config) -> dict: @@ -327,9 +332,9 @@ def test_emitted_keys_are_the_model_state_dict(checkpoint): def test_expert_quant_tag_follows_the_config(checkpoint): name, folder, _raw = checkpoint config = parse_config(cached_load_hf_config(folder)) - expected = {"bf16": "none", "fp8_block": "fp8_block", "ct_nvfp4_dense": "none"}.get(name, "nvfp4") + expected = {"bf16": "none", "fp8_block": "fp8_block", "ct_block_experts": "fp8_block", "ct_nvfp4_dense": "none"}.get(name, "nvfp4") assert config.expert_quant == expected - assert config.weight_block_size == ((128, 128) if name == "fp8_block" else None) + assert config.weight_block_size == ((128, 128) if expected == "fp8_block" else None) def _slices(fused: torch.Tensor, parts: list[torch.Tensor]) -> list[torch.Tensor]: @@ -372,7 +377,7 @@ def test_bf16_stacked_experts_pass_through_only_when_asked(checkpoint): def test_block_fp8_fuses_weight_and_scale_per_kind(checkpoint): name, folder, raw = checkpoint - if name not in ("fp8_block", "ct_block_moe"): + if name not in ("fp8_block", "ct_block_moe", "ct_block_experts"): pytest.skip("block-fp8 layouts only") loaded = _load(folder) scale = "weight_scale_inv" if name == "fp8_block" else "weight_scale" @@ -498,6 +503,23 @@ def test_nvfp4_expert_pieces_read_either_dialect_from_a_single_file(checkpoint): assert torch.equal(piece["gate_global"].reshape(-1), (1.0 / raw[f"{base}.weight_global_scale"]).to(torch.float16)) +@pytest.mark.parametrize("parallel", [False, True]) +def test_block_fp8_expert_pieces_read_either_dialect(checkpoint, parallel): + """The block-fp8 expert reader takes the scale's name from the dialect: ``weight_scale_inv`` (HF fp8) or ``weight_scale`` (llm-compressor).""" + name, folder, raw = checkpoint + if name not in ("fp8_block", "ct_block_experts"): + pytest.skip("block-fp8 expert layouts only") + _install(folder) + config = parse_config(cached_load_hf_config(folder)) + pieces = list(iter_expert_pieces(folder, config, QuantKind.FP8_BLOCK, parallel=parallel)) + assert len(pieces) == 2 * E + layer, e0, e1, piece = next(p for p in pieces if p[0] == 1 and p[1] == 2) + base = f"{LM}.layers.1.mlp.experts.2" + scale = "weight_scale_inv" if name == "fp8_block" else "weight_scale" + assert _same(piece["gate"][0], raw[f"{base}.gate_proj.weight"]) + assert torch.equal(piece["down_scale"][0].float(), raw[f"{base}.down_proj.{scale}"].float()) + + # --------------------------------------------------------------------------- the family's unquantized modules