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