From d68a08eddc540cd3872543440c1fd08464a891a4 Mon Sep 17 00:00:00 2001 From: gberasmus87 Date: Tue, 1 Sep 2026 15:25:16 +1200 Subject: [PATCH 1/3] qwen4_exp: load modelopt MIXED_PRECISION (NVFP4 experts + block-FP8 dense) checkpoints --- python/freetoken/models/qwen4_exp/config.py | 15 ++++++++++++ python/freetoken/models/qwen4_exp/weight.py | 26 +++++++++++++++++++-- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index bb5d1dff5..03648bfc3 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -163,6 +163,21 @@ def parse_config(hf_config: Any) -> ModelConfig: assert bs == (128, 128), f"only 128x128 block-fp8 is supported, got {bs}" expert_quant = "fp8_block" attn_quant = dense_quant = lm_head_quant = "none" + elif algo == "mixed_precision": + # modelopt MIXED_PRECISION: the quant algo is declared per module in + # ``quantized_layers`` rather than once at the top level. The community + # NVFP4-FP8 build of Qwen3.8-Flash-Next quantizes the routed experts to NVFP4 + # (read natively by the offload cache) and the dense attn/GDN projections to + # 128x128 block-FP8; the block-FP8 dense weights are dequantized to bf16 at load + # (see weight.py ``_load_maybe_block_fp8``), so every non-expert module is bf16. + quantized = get("quantized_layers") or {} + experts_nvfp4 = any( + ".mlp.experts" in str(module) + and str((spec or {}).get("quant_algo", "")).upper() == "NVFP4" + for module, spec in quantized.items() + ) + expert_quant = "nvfp4" if experts_nvfp4 else "none" + attn_quant = dense_quant = lm_head_quant = "none" else: is_fp4 = "fp4" in algo ignore = list(get("ignore") or []) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index f8d2a7494..66e6036d2 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -46,7 +46,8 @@ desc="Qwen3.8-Flash-Next NVFP4 experts", ) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. -_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +# ``.weight_scale_inv`` is the 128x128 block-FP8 reciprocal scale (see _load_maybe_block_fp8). +_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".weight_scale_inv", ".input_scale") # The n-gram table itself: too big for the dense state dict, loaded by load_ple_table. _PLE_TABLE_INFIX = ".ple.ple_embedding.ngram_embedding." @@ -137,6 +138,26 @@ def _try_fuse( return None +def _load_maybe_block_fp8(f, raw_name: str, keyset: set[str]) -> torch.Tensor: + """Load ``raw_name``, dequantizing 128x128 block-FP8 to bf16 when a sibling + ``.weight_scale_inv`` is present in the same shard; pass plain bf16 through unchanged. + + The official modelopt checkpoint keeps the dense attn/GDN/HC/PLE projections bf16 (they are + on the quant ``ignore`` list), but some community requants -- e.g. the lovedheart NVFP4-FP8 + build that fits Qwen3.8-Flash-Next on a 24 GB card -- store those dense weights as block-FP8. + Without this they reach ``_try_fuse`` as fp8 and crash on the fp8+bf16 ``torch.cat``.""" + tensor = f.get_tensor(raw_name) + if raw_name.endswith(".weight"): + base = raw_name[: -len(".weight")] + if base + ".weight_scale_inv" in keyset: + from freetoken.kernel.triton.fp8_block_linear import dequant_block_fp8 + + return dequant_block_fp8( + tensor, f.get_tensor(base + ".weight_scale_inv") + ).to(torch.bfloat16) + return tensor + + def iter_weights( model_path: str, device: torch.device, @@ -169,11 +190,12 @@ def iter_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(): name = _rename(raw_name) if name is None: continue - tensor = f.get_tensor(raw_name) + tensor = _load_maybe_block_fp8(f, raw_name, keyset) fused = _try_fuse(name, tensor, fuse_buf) if fused is not None: if fused != (): # () means buffered, not yet complete From da06515da7d0f4bad4a0f1a62fe60733dcf6b943 Mon Sep 17 00:00:00 2001 From: gberasmus87 Date: Fri, 4 Sep 2026 22:52:07 +0000 Subject: [PATCH 2/3] qwen4_exp: serve the block-FP8 dense projections natively #320 loads modelopt MIXED_PRECISION checkpoints by dequantizing the block-FP8 dense projections to bf16 at load, which doubles the bytes read on every decode step. Keep them quantized instead and let the existing Fp8Block linears consume them: 4.98 GiB -> 2.49 GiB of dense weights on the modelopt build of Qwen3.8-Flash-Next (156 FP8_PB_WO tensors, 2.67G elements). The checkpoint declares this per module - 48 .mlp.experts as NVFP4 and 156 attn/GDN projections as FP8_PB_WO - so config.py now reports attn_quant="fp8_block" independently of expert_quant, and gdn.py takes the block-fp8 path when either says so. It previously keyed off expert_quant alone, so a checkpoint with NVFP4 experts and block-FP8 dense never reached it. quant_linear.py's factories widen the same way, and qwen4_exp attention builds qkv_proj/o_proj through them instead of hardcoding the bf16 classes. The loader stops dequantizing when the declaration is present, keeps the weight_scale_inv tensors, and swaps in a fusion table matching the modules the model actually builds: the four-way in_proj fusion splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM - the split gdn.py already implements for block-fp8, matching sglang/vLLM - and each fp8 group fuses its scale on the same axis as its weight. That split is what unblocks the bulk of this: b|a are bf16, so the old four-way cat mixed dtypes and forced the dequant. Every fp8 part is a whole number of 128-row blocks (10240/6144 and 12288/512/512), so the per-block scales concatenate exactly alongside the rows they describe. A checkpoint carrying weight_scale_inv WITHOUT declaring FP8_PB_WO still takes the #320 dequant path, so builds that quantize the dense side but describe it differently keep working unchanged. Verified on the modelopt checkpoint without a GPU: parse_config yields nvfp4/fp8_block, and "none" when the declaration is removed; iter_weights emits in_proj_qkvz [16384,2560] fp8 + scale [128,20], in_proj_ba [96,2560] bf16, qkv_proj [13312,2560] fp8 + scale [104,20], and fp8 out_proj/o_proj - the exact buffers Fp8BlockLinear declares. Serving numbers to follow. --- python/freetoken/models/quant_linear.py | 7 +- .../freetoken/models/qwen4_exp/attention.py | 9 +- python/freetoken/models/qwen4_exp/config.py | 20 +++- python/freetoken/models/qwen4_exp/gdn.py | 5 +- python/freetoken/models/qwen4_exp/weight.py | 92 +++++++++++++++++-- 5 files changed, 114 insertions(+), 19 deletions(-) diff --git a/python/freetoken/models/quant_linear.py b/python/freetoken/models/quant_linear.py index 6ba67a9ee..a937541bd 100644 --- a/python/freetoken/models/quant_linear.py +++ b/python/freetoken/models/quant_linear.py @@ -2,7 +2,8 @@ dense projections (qwen3_5_moe, muse_glimmer). Maps the model's quant config (``expert_quant`` for the dense MLP / shared-expert path, -``attn_quant`` for attention + GatedDeltaNet projections) to the right ``BaseOP`` linear: +``attn_quant`` for attention + GatedDeltaNet projections; block-fp8 may be declared by +either, since a modelopt MIXED_PRECISION checkpoint block-quantizes only the dense side) to the right ``BaseOP`` linear: block-FP8, per-tensor-FP8 and NVFP4 implementations live under ``freetoken.kernel.triton``; the bf16 fallback is the framework's TP-aware ``freetoken.layers``. Only the *dispatch* (config -> layer class) lives here. @@ -14,7 +15,7 @@ def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int, output_sizes: list[int], has_bias: bool = False): """Column-merged linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16.""" - if expert_quant == "fp8_block": + if "fp8_block" in (expert_quant, attn_quant): from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged return Fp8BlockColMerged(in_f, output_sizes, has_bias) @@ -34,7 +35,7 @@ def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int, def make_replicated_quant(expert_quant: str, attn_quant: str, in_f: int, out_f: int, has_bias: bool = False): """Replicated linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16.""" - if expert_quant == "fp8_block": + if "fp8_block" in (expert_quant, attn_quant): from freetoken.kernel.triton.fp8_block_linear import Fp8BlockLinear return Fp8BlockLinear(in_f, out_f, has_bias) diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index d6ab2867a..ba6279b61 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -20,6 +20,7 @@ import torch from freetoken.core import get_global_ctx from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.models.quant_linear import make_col_merged, make_replicated from freetoken.layers.rotary import get_rope from freetoken.utils import nvtx_annotate @@ -121,10 +122,10 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: self.qo_attn_dim = self.num_q * self.head_dim self.kv_attn_dim = self.num_kv * self.head_dim self._qkv_split = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] - self.qkv_proj = LinearColParallelMerged( - config.hidden_size, self._qkv_split, has_bias=False - ) - self.o_proj = LinearReplicated(self.qo_attn_dim, config.hidden_size, has_bias=False) + # q|k|v are all quantized together (or all bf16), so the merged GEMM stays a + # single kernel; a modelopt MIXED_PRECISION checkpoint declares them FP8_PB_WO. + self.qkv_proj = make_col_merged(config, config.hidden_size, self._qkv_split) + self.o_proj = make_replicated(config, self.qo_attn_dim, config.hidden_size) self.q_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) rotary = config.rotary_config diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 03648bfc3..98cde5654 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -120,6 +120,10 @@ def _layer_types(text: Any) -> list[str]: ] +# modelopt spellings for 128x128 per-block, weight-only FP8 on the dense modules. +_FP8_BLOCK_ALGOS = frozenset({"FP8_PB_WO", "FP8_BLOCK"}) + + def parse_config(hf_config: Any) -> ModelConfig: text = getattr(hf_config, "text_config", hf_config) @@ -168,16 +172,26 @@ def parse_config(hf_config: Any) -> ModelConfig: # ``quantized_layers`` rather than once at the top level. The community # NVFP4-FP8 build of Qwen3.8-Flash-Next quantizes the routed experts to NVFP4 # (read natively by the offload cache) and the dense attn/GDN projections to - # 128x128 block-FP8; the block-FP8 dense weights are dequantized to bf16 at load - # (see weight.py ``_load_maybe_block_fp8``), so every non-expert module is bf16. + # 128x128 block-FP8, declared per module as ``FP8_PB_WO``. quantized = get("quantized_layers") or {} experts_nvfp4 = any( ".mlp.experts" in str(module) and str((spec or {}).get("quant_algo", "")).upper() == "NVFP4" for module, spec in quantized.items() ) + # The same map declares the dense attn/GDN projections as FP8_PB_WO + # (per-block, weight-only FP8 with a ``weight_scale_inv`` sibling). Serve + # them natively instead of dequantizing at load: the four-way in_proj fusion + # splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM (see gdn.py), which + # halves the dense bytes read on every decode step. + dense_block_fp8 = any( + ".mlp.experts" not in str(module) + and str((spec or {}).get("quant_algo", "")).upper() in _FP8_BLOCK_ALGOS + for module, spec in quantized.items() + ) expert_quant = "nvfp4" if experts_nvfp4 else "none" - attn_quant = dense_quant = lm_head_quant = "none" + attn_quant = "fp8_block" if dense_block_fp8 else "none" + dense_quant = lm_head_quant = "none" else: is_fp4 = "fp4" in algo ignore = list(get("ignore") or []) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index 69838153f..1f3e6e44b 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -83,7 +83,10 @@ def __init__( # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). - self._block_fp8 = expert_quant == "fp8_block" + # Block-fp8 dense can come from either a wholly block-fp8 checkpoint + # (expert_quant) or a modelopt MIXED_PRECISION one where only the dense + # attn/GDN projections are FP8_PB_WO while the experts are NVFP4 (attn_quant). + self._block_fp8 = "fp8_block" in (expert_quant, attn_quant) self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._fp8 = self._block_fp8 or self._pertensor_fp8 diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 66e6036d2..70e5097cd 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -99,15 +99,20 @@ } -def _rename(raw_name: str) -> str | None: - """Checkpoint key -> FreeToken state-dict key, or None to skip.""" +def _rename(raw_name: str, keep_scale_inv: bool = False) -> str | None: + """Checkpoint key -> FreeToken state-dict key, or None to skip. + + ``keep_scale_inv`` retains the block-FP8 ``weight_scale_inv`` tensors, which the + fp8 linears need alongside their weight; they are dropped otherwise.""" if raw_name.startswith(("mtp.", "model.visual.", "visual.")): return None if _PLE_TABLE_INFIX in raw_name: return None # n-gram table + its scale: load_ple_table if _EXPERT_RE.search(raw_name): return None # routed experts: offload source banks - if raw_name.endswith(_SCALE_SUFFIXES): + if raw_name.endswith(_SCALE_SUFFIXES) and not ( + keep_scale_inv and raw_name.endswith(".weight_scale_inv") + ): return None if raw_name.startswith("model.language_model."): return "model." + raw_name[len("model.language_model.") :] @@ -117,10 +122,11 @@ def _rename(raw_name: str) -> str | None: def _try_fuse( - name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]] + name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]], + table: dict[str, tuple[tuple[str, ...], int]] | None = None, ) -> tuple[str, torch.Tensor] | tuple[()] | None: """Buffer a fusion part; return the merged ``(name, tensor)`` once all parts arrive, ``()`` while incomplete, ``None`` if ``name`` is not a fusion part.""" - for fused_suffix, (parts, pad_to) in _FUSIONS.items(): + for fused_suffix, (parts, pad_to) in (table or _FUSIONS).items(): for idx, part in enumerate(parts): if not name.endswith(part): continue @@ -158,6 +164,69 @@ def _load_maybe_block_fp8(f, raw_name: str, keyset: set[str]) -> torch.Tensor: return tensor +# modelopt spellings for 128x128 per-block, weight-only FP8 on the dense modules. +_FP8_BLOCK_ALGOS = frozenset({"FP8_PB_WO", "FP8_BLOCK"}) + +# Serving the dense side natively as block-FP8 changes which buffers the model expects: +# the four-way in_proj fusion splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM +# (see gdn.py), and each fp8 group fuses its ``weight_scale_inv`` on the same axis as its +# ``weight``. Every fp8 part is a whole number of 128-row blocks, so the per-block scales +# concatenate exactly alongside the rows they describe. +_BLOCK_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", + ), +} +_BLOCK_BF16_FUSE: dict[str, tuple[str, ...]] = { + ".linear_attn.in_proj_ba": (".linear_attn.in_proj_b", ".linear_attn.in_proj_a"), +} +_BLOCK_FP8_KINDS = (".weight", ".weight_scale_inv") + + +def _block_fp8_fusions() -> dict[str, tuple[tuple[str, ...], int]]: + """``_FUSIONS`` with the two attention groups replaced by their block-FP8 split.""" + table = { + key: val + for key, val in _FUSIONS.items() + if key not in (".self_attn.qkv_proj.weight", ".linear_attn.in_proj.weight") + } + for fused, parts in _BLOCK_FP8_FUSE.items(): + for kind in _BLOCK_FP8_KINDS: + table[fused + kind] = (tuple(part + kind for part in parts), 0) + for fused, parts in _BLOCK_BF16_FUSE.items(): + table[fused + ".weight"] = (tuple(part + ".weight" for part in parts), 0) + return table + + +_FUSIONS_BLOCK_FP8 = _block_fp8_fusions() + + +def _dense_is_block_fp8(model_path: str) -> bool: + """True when the checkpoint DECLARES its dense (non-expert) modules per-block + weight-only FP8 -- exactly when config.py sets ``attn_quant="fp8_block"``. + + Both sides read the same declaration, so the buffers emitted here always match the + modules the model built. A checkpoint carrying ``weight_scale_inv`` WITHOUT declaring + it falls through to ``_load_maybe_block_fp8`` and is dequantized to bf16 as before. + """ + try: + with open(os.path.join(model_path, "config.json"), encoding="utf-8") as fh: + quant = json.load(fh).get("quantization_config") or {} + except (OSError, ValueError): + return False + algo = str(quant.get("quant_algo") or quant.get("quant_method") or "").lower() + if algo != "mixed_precision": + return False + return any( + ".mlp.experts" not in str(module) + and str((spec or {}).get("quant_algo", "")).upper() in _FP8_BLOCK_ALGOS + for module, spec in (quant.get("quantized_layers") or {}).items() + ) + + def iter_weights( model_path: str, device: torch.device, @@ -183,6 +252,9 @@ def iter_weights( if not include_non_moe: return + # Declared block-FP8 dense is served natively; anything else keeps the dequant path. + block_fp8 = _dense_is_block_fp8(model_path) + fusions = _FUSIONS_BLOCK_FP8 if block_fp8 else _FUSIONS fuse_buf: dict[str, dict[int, torch.Tensor]] = {} for file in tqdm( iter_weight_files(model_path), @@ -192,11 +264,15 @@ def iter_weights( with safetensors.safe_open(file, framework="pt", device=str(device)) as f: keyset = set(f.keys()) for raw_name in f.keys(): - name = _rename(raw_name) + name = _rename(raw_name, keep_scale_inv=block_fp8) if name is None: continue - tensor = _load_maybe_block_fp8(f, raw_name, keyset) - fused = _try_fuse(name, tensor, fuse_buf) + tensor = ( + f.get_tensor(raw_name) + if block_fp8 + else _load_maybe_block_fp8(f, raw_name, keyset) + ) + fused = _try_fuse(name, tensor, fuse_buf, fusions) if fused is not None: if fused != (): # () means buffered, not yet complete yield fused From 72773b04dd4ce76ea6e60b2c0e455b83ad4dfec8 Mon Sep 17 00:00:00 2001 From: gberasmus87 Date: Sat, 5 Sep 2026 03:13:48 +0000 Subject: [PATCH 3/3] qwen4_exp: keep the TP-aware linears, and resolve the dense mode once Both from @gdevenyi's review on 2 x RTX 6000 Ada, where this is carried on a deploy branch alongside #385 (TP). 1. Attention routed its bf16 fallback through the quantized factories too, which swaps in their generic fallback and drops the tensor-parallel classes #385 needs (per-rank local_output_sizes, row-parallel o_proj). That is exactly the path a rank takes under TP>1, since the block-FP8 linears have no parallel variant. The factories are now used only on the fp8_block branch; every other case keeps LinearColParallelMerged / LinearReplicated as before. 2. config.parse_config and weight._dense_is_block_fp8 read the same declaration through two independent code paths, each with its own copy of _FP8_BLOCK_ALGOS. That is safe only while they cannot disagree, and they can: a rank downgrading under TP>1 must have the modules it BUILDS and the buffers it LOADS downgrade together, or the buffers will not match. Both now resolve through one helper, config.dense_quant_mode, which owns the declaration test and the TP downgrade. The duplicate constant is gone. It reads TP through try_get_tp_info, not get_tp_info: Engine.__init__ sets TP info as its first statement so a rank always knows its size by the time this matters, but config parsing also happens with no engine at all (checkpoint conversion, tooling, tests) where get_tp_info raises. Verified on the modelopt checkpoint: parse_config still yields nvfp4/fp8_block; the two sides agree at TP=1 (both fp8_block) and at TP=2 (both downgraded); attention builds LinearColParallelMerged/LinearReplicated under bf16 and Fp8BlockColMerged/Fp8BlockLinear under fp8_block. tests/models/qwen4_exp/test_config.py + test_weight.py: 30 passed. The whole qwen4_exp suite reports 47 failed / 46 passed / 50 skipped both at the merge-base and with these fixes - identical sets, no regressions. Those failures are pre-existing and are an artefact of this box rather than the code: its single 24 GB card is 23.6 GB occupied serving a model, so the GPU-dependent tests cannot allocate. I have not been able to run them on a free card. --- .../freetoken/models/qwen4_exp/attention.py | 16 ++++++- python/freetoken/models/qwen4_exp/config.py | 47 ++++++++++++++++--- python/freetoken/models/qwen4_exp/weight.py | 28 +++++------ 3 files changed, 66 insertions(+), 25 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index ba6279b61..9bd5a154c 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -124,8 +124,20 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: self._qkv_split = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] # q|k|v are all quantized together (or all bf16), so the merged GEMM stays a # single kernel; a modelopt MIXED_PRECISION checkpoint declares them FP8_PB_WO. - self.qkv_proj = make_col_merged(config, config.hidden_size, self._qkv_split) - self.o_proj = make_replicated(config, self.qo_attn_dim, config.hidden_size) + # + # The quantized factories are used ONLY on the block-FP8 path. Routing the bf16 + # case through them too would swap in the factory's generic fallback and drop the + # tensor-parallel classes this model needs (per-rank ``local_output_sizes``, a + # row-parallel ``o_proj``), which is exactly the bf16 path a rank falls back to + # under TP>1 - where the block-FP8 linears have no parallel variant. + if getattr(config, "attn_quant", "none") == "fp8_block": + self.qkv_proj = make_col_merged(config, config.hidden_size, self._qkv_split) + self.o_proj = make_replicated(config, self.qo_attn_dim, config.hidden_size) + else: + self.qkv_proj = LinearColParallelMerged( + config.hidden_size, self._qkv_split, has_bias=False + ) + self.o_proj = LinearReplicated(self.qo_attn_dim, config.hidden_size, has_bias=False) self.q_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) rotary = config.rotary_config diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 98cde5654..c24712bc2 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -6,6 +6,7 @@ import torch +from freetoken.distributed import try_get_tp_info from freetoken.models.config import ( FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, @@ -124,6 +125,41 @@ def _layer_types(text: Any) -> list[str]: _FP8_BLOCK_ALGOS = frozenset({"FP8_PB_WO", "FP8_BLOCK"}) +def dense_quant_mode(algo: str, quantized_layers: Any) -> str: + """The quantization mode the dense (non-expert) projections will actually be SERVED in. + + Single source of truth for the two sides that must agree: :func:`parse_config`, which + decides the modules the model BUILDS, and the weight loader, which decides the buffers + it EMITS. They previously derived this independently from the same declaration - safe + only while they cannot disagree, and they can: the block-FP8 linears have no + tensor-parallel variant, so a rank running under TP>1 has to fall back to bf16. If only + one side knew that, the loaded buffers would not match the built modules. + + Returns ``"fp8_block"`` only when the checkpoint declares per-block weight-only FP8 on a + non-expert module AND this rank can serve it; ``"none"`` otherwise (bf16, via the + dequantize-at-load path). A checkpoint carrying ``weight_scale_inv`` without declaring + the algo is "none" here and keeps the pre-existing dequant behaviour. + """ + if str(algo or "").lower() != "mixed_precision": + return "none" + declared = any( + ".mlp.experts" not in str(module) + and str((spec or {}).get("quant_algo", "")).upper() in _FP8_BLOCK_ALGOS + for module, spec in (quantized_layers or {}).items() + ) + if not declared: + return "none" + # Resolved here, once, so both sides downgrade together. ``Engine.__init__`` sets TP + # info as its very first statement, before the model config or any weight is built, so + # a rank always knows its size by the time this matters. try_get_tp_info is used rather + # than get_tp_info because config parsing also happens with no engine at all (checkpoint + # conversion, tooling, tests), where get_tp_info raises; unset means a single rank. + tp = try_get_tp_info() + if tp is not None and tp.size > 1: + return "none" + return "fp8_block" + + def parse_config(hf_config: Any) -> ModelConfig: text = getattr(hf_config, "text_config", hf_config) @@ -183,14 +219,11 @@ def parse_config(hf_config: Any) -> ModelConfig: # (per-block, weight-only FP8 with a ``weight_scale_inv`` sibling). Serve # them natively instead of dequantizing at load: the four-way in_proj fusion # splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM (see gdn.py), which - # halves the dense bytes read on every decode step. - dense_block_fp8 = any( - ".mlp.experts" not in str(module) - and str((spec or {}).get("quant_algo", "")).upper() in _FP8_BLOCK_ALGOS - for module, spec in quantized.items() - ) + # halves the dense bytes read on every decode step. Resolved through + # dense_quant_mode so the loader reaches the same answer, TP downgrade + # included - see that function. expert_quant = "nvfp4" if experts_nvfp4 else "none" - attn_quant = "fp8_block" if dense_block_fp8 else "none" + attn_quant = dense_quant_mode(algo, quantized) dense_quant = lm_head_quant = "none" else: is_fp4 = "fp4" in algo diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 70e5097cd..e1f2b22ed 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -22,6 +22,7 @@ import torch from freetoken.distributed import get_tp_info from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.qwen4_exp.config import dense_quant_mode from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, load_nvfp4_expert_source_banks, @@ -164,9 +165,6 @@ def _load_maybe_block_fp8(f, raw_name: str, keyset: set[str]) -> torch.Tensor: return tensor -# modelopt spellings for 128x128 per-block, weight-only FP8 on the dense modules. -_FP8_BLOCK_ALGOS = frozenset({"FP8_PB_WO", "FP8_BLOCK"}) - # Serving the dense side natively as block-FP8 changes which buffers the model expects: # the four-way in_proj fusion splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM # (see gdn.py), and each fp8 group fuses its ``weight_scale_inv`` on the same axis as its @@ -205,26 +203,24 @@ def _block_fp8_fusions() -> dict[str, tuple[tuple[str, ...], int]]: def _dense_is_block_fp8(model_path: str) -> bool: - """True when the checkpoint DECLARES its dense (non-expert) modules per-block - weight-only FP8 -- exactly when config.py sets ``attn_quant="fp8_block"``. + """True when the dense (non-expert) modules will be SERVED as per-block weight-only + FP8 -- exactly when config.py sets ``attn_quant="fp8_block"``. + + The decision itself lives in :func:`~freetoken.models.qwen4_exp.config.dense_quant_mode` + and is only read here, so the buffers this loader emits cannot disagree with the modules + the model built. That matters under TP: the block-FP8 linears have no parallel variant, + so a rank at TP>1 downgrades to bf16, and both sides have to downgrade together. - Both sides read the same declaration, so the buffers emitted here always match the - modules the model built. A checkpoint carrying ``weight_scale_inv`` WITHOUT declaring - it falls through to ``_load_maybe_block_fp8`` and is dequantized to bf16 as before. + A checkpoint carrying ``weight_scale_inv`` WITHOUT declaring the algo is not block-FP8 + here either; it falls through to ``_load_maybe_block_fp8`` and is dequantized as before. """ try: with open(os.path.join(model_path, "config.json"), encoding="utf-8") as fh: quant = json.load(fh).get("quantization_config") or {} except (OSError, ValueError): return False - algo = str(quant.get("quant_algo") or quant.get("quant_method") or "").lower() - if algo != "mixed_precision": - return False - return any( - ".mlp.experts" not in str(module) - and str((spec or {}).get("quant_algo", "")).upper() in _FP8_BLOCK_ALGOS - for module, spec in (quant.get("quantized_layers") or {}).items() - ) + algo = quant.get("quant_algo") or quant.get("quant_method") or "" + return dense_quant_mode(algo, quant.get("quantized_layers")) == "fp8_block" def iter_weights(