From a091aa056d4a236a5729ea61acd87617e4102702 Mon Sep 17 00:00:00 2001 From: Chris Qian Date: Wed, 26 Aug 2026 12:17:11 +0800 Subject: [PATCH 1/3] qwen3_5_moe: load mixed-precision compressed-tensors checkpoints (unsloth NVFP4+FP8) unsloth's dynamic per-module quant exports (e.g. unsloth/Qwen3.8-27B-NVFP4) store dense Qwen3.x checkpoints as per-module mixed precision: FP8 attention / GDN output linears (per-row scale), NVFP4 dense-MLP layers, FP8 dense-MLP layers, bf16 in_proj_b/a and norms. The loader crashed on them (fp8/bf16 promotion in the bf16 fusion, missing packed weights for natively-built linears) or materialized everything to bf16 (54 GB on a 27B -- no launch parameter fits a 32 GB card). Keep every dense linear in the storage the checkpoint actually uses, sniffed from model.safetensors.index.json: - config: _compressed_linear_storage() now reports per-module storage (attention nvfp4/fp8/none, per-layer dense-MLP overrides, fp8 lm_head); ModelConfig gains dense_mlp_storage (per-layer override map) and routes attn_quant/lm_head_quant to native fp8_pertensor when the export says so - moe: _SharedExpert builds per-layer native linears (NVFP4 W4A16 / FP8 W8A16 / bf16) from the override map; layer_id threaded through the dense MLP - model: lm_head built as native FP8 (Fp8PerTensorLinear) when the checkpoint stores it fp8 (halves the ~2.5 GB bf16 lm_head) - weight: the dense pass keeps fp8 parts native (q/k/v -> qkv_proj, in_proj_qkv/z -> in_proj_qkvz, dense gate/up -> gate_up_proj fusions; o_proj/out_proj/down_proj/lm_head singletons, per-row fp32 scales); a part buffered into the fp8 fusion never also enters the bf16 buffer (incomplete fusion assert); fp8 dequant (for the unscaled remainder) is a bf16 broadcast multiply (no fp32 copy); ShardReader gains has() for sibling-scale lookups Verified on unsloth/Qwen3.8-27B-NVFP4 (RTX 5090 D, 32 GB): weights resident at ~21.8 GB native (vs 54 GB bf16), CUDA graph capture, and live chat-completion requests all succeed; official dense-NVFP4 and routed-MoE layouts keep their native assumptions (sniffer fallbacks) and are unchanged. Tests: tests/models/test_qwen3_5_moe_config.py (8) + test_qwen3_5_moe_weight.py (13) -- storage sniffing, per-layer construction gates, fp8 native fusions, per-row/block scale dequant semantics. --- python/freetoken/models/config.py | 5 + python/freetoken/models/loader.py | 3 + python/freetoken/models/qwen3_5_moe/config.py | 89 +++++++- python/freetoken/models/qwen3_5_moe/model.py | 12 +- python/freetoken/models/qwen3_5_moe/moe.py | 44 +++- python/freetoken/models/qwen3_5_moe/weight.py | 131 +++++++++-- tests/models/test_qwen3_5_moe_config.py | 167 ++++++++++++++ tests/models/test_qwen3_5_moe_weight.py | 211 ++++++++++++++++++ 8 files changed, 637 insertions(+), 25 deletions(-) create mode 100644 tests/models/test_qwen3_5_moe_config.py create mode 100644 tests/models/test_qwen3_5_moe_weight.py diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1f..5a206520 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -243,6 +243,11 @@ class ModelConfig: # it bf16. Separate from dense_quant because only some NVFP4 checkpoints quantize lm_head # (modelopt MIXED_PRECISION does; pure NVFP4 leaves it bf16). lm_head_quant: str = "none" + # Per-layer override of the dense-MLP storage for mixed compressed-tensors exports + # (unsloth: most layers packed NVFP4, a few FP8 layers). Keys are layer ids; the value + # replaces dense_quant for that layer's dense MLP at construction ("nvfp4" native W4A16 + # / "bf16" dequant-at-load). None = every dense-MLP layer follows dense_quant. + dense_mlp_storage: dict[int, str] | None = None shared_expert_intermediate_size: int = 0 use_qk_norm: bool = False # ----- DeepSeek/GLM-style MoE extensions (default keeps other models intact) ----- diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 49419364..5e14c7fe 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -182,6 +182,9 @@ def files(self) -> list[str]: def names_in(self, file: str) -> list[str]: return [name for name, shard in self._map.items() if shard == file] + def has(self, name: str) -> bool: + return name in self._map + def get_tensor(self, name: str) -> torch.Tensor: import safetensors diff --git a/python/freetoken/models/qwen3_5_moe/config.py b/python/freetoken/models/qwen3_5_moe/config.py index dc16cfff..b896aac4 100644 --- a/python/freetoken/models/qwen3_5_moe/config.py +++ b/python/freetoken/models/qwen3_5_moe/config.py @@ -68,6 +68,77 @@ def _expert_quant(hf_config: Any) -> str: _compressed_tensors_nvfp4 = detect_compressed_tensors_nvfp4 +def _compressed_linear_storage( + hf_config: Any, +) -> tuple[str, str, dict[int, str] | None, bool]: + """Per-module dense-linear storage sniffed from the checkpoint's + ``model.safetensors.index.json`` weight_map: + ``(attention storage, dense-MLP fallback, per-layer dense-MLP overrides, lm_head + is fp8)``. + + The attention storage is ``"nvfp4"`` (packed q/k/v/o + GDN out_proj linears), + ``"fp8"`` (weight + scale siblings, e.g. unsloth's mixed export), or ``"none"`` + (bf16). The fallback is ``"nvfp4"`` when at least one dense-MLP layer is packed, + else ``"none"``. The override map gives the storage of dense-MLP layers that are + not packed NVFP4 (``"fp8"`` keeps them native W8A16; ``"bf16"`` dequantizes at + load); ``None`` means every dense-MLP layer follows the fallback. When the index + is unavailable (a hub id before download, or a single-file checkpoint) assume + the official dense-NVFP4 layout (all packed, no overrides, bf16 lm_head). + """ + import json + import os + import re + + path = getattr(hf_config, "name_or_path", None) + if not isinstance(path, str) or not os.path.isdir(path): + return "nvfp4", "nvfp4", None, False + index = os.path.join(path, "model.safetensors.index.json") + if not os.path.isfile(index): + return "nvfp4", "nvfp4", None, False + with open(index, encoding="utf-8") as f: + weight_map = json.load(f).get("weight_map", {}) + keys = set(weight_map) + attn_suffix_re = re.compile(r"\.(?:self_attn\.(?:q|k|v|o)_proj|linear_attn\.out_proj)\.") + attn_bases = set() + for k in keys: + if attn_suffix_re.search(k) and k.endswith((".weight", ".weight_packed")): + attn_bases.add(k.rsplit(".", 1)[0]) + attn_packed = any(b + ".weight_packed" in keys for b in attn_bases) + attn_fp8 = any(b + ".weight_scale" in keys for b in attn_bases) + attn = "nvfp4" if attn_packed else ("fp8" if attn_fp8 else "none") + # Per-layer dense-MLP storage (order-independent): packed -> nvfp4, any scale + # sibling -> fp8 (native W8A16), bare weight -> bf16. lm_head: fp8 when it ships a + # scale sibling. + mlp_key_re = re.compile( + r"\.layers\.(\d+)\.mlp\.(?:gate|up|down)_proj\.(weight_packed|weight_scale|weight)$" + ) + packed_layers: set[int] = set() + scaled_layers: set[int] = set() + plain_layers: set[int] = set() + for k in keys: + mm = mlp_key_re.search(k) + if mm: + layer = int(mm.group(1)) + suffix = mm.group(2) + if suffix == "weight_packed": + packed_layers.add(layer) + elif suffix == "weight_scale": + scaled_layers.add(layer) + else: + plain_layers.add(layer) + lmhead_fp8 = any(k == "lm_head.weight_scale" or k.endswith(".lm_head.weight_scale") for k in keys) + # No mlp.{gate,up,down}_proj keys at all: routed-expert MoE naming (shared_expert / + # experts) -- keep the native assumption with no per-layer overrides. + if not (packed_layers or scaled_layers or plain_layers): + return attn, "nvfp4", None, lmhead_fp8 + all_layers = packed_layers | scaled_layers | plain_layers + overrides = { + layer: ("fp8" if layer in scaled_layers else "bf16") + for layer in all_layers - packed_layers + } + return attn, "nvfp4" if packed_layers else "none", overrides or None, lmhead_fp8 + + 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 @@ -182,10 +253,21 @@ def parse_config(hf_config: Any) -> ModelConfig: # 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). + # Mixed compressed-tensors exports (unsloth's dynamic per-module quant) store the + # attention projections as weight-only FP8 (per-row scale) and some dense-MLP layers + # as FP8: keep each group in its native storage (W8A16 attention, W4A16 packed MLP + # layers, bf16 dequant for the rest) instead of dequantizing everything to bf16. + dense_mlp_storage: dict[int, str] | None = None if _compressed_tensors_nvfp4(hf_config): - attn_quant = "nvfp4" - dense_quant = "nvfp4" - lm_head_quant = "none" + attn_storage, dense_quant, dense_mlp_storage, lmhead_fp8 = _compressed_linear_storage(hf_config) + attn_quant = ( + "fp8_pertensor" if attn_storage == "fp8" + else "nvfp4" if attn_storage == "nvfp4" + else "none" + ) + # unsloth's mixed export stores lm_head as weight-only FP8 (per-row scale): keep + # it native W8A16 (halves the ~2.5 GB bf16 lm_head); official layouts are bf16. + lm_head_quant = "fp8_pertensor" if lmhead_fp8 else "none" # 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. @@ -257,6 +339,7 @@ def parse_config(hf_config: Any) -> ModelConfig: attn_quant=attn_quant, dense_quant=dense_quant, lm_head_quant=lm_head_quant, + dense_mlp_storage=dense_mlp_storage, ) diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24..866194e4 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -49,7 +49,7 @@ def __init__(self, config: ModelConfig, layer_id: int): self.self_attn = Qwen3_5Attention(config, layer_id) # Dense variants (num_experts==0, e.g. Qwen3.6-27B) use a plain SwiGLU MLP instead of # the routed MoE block; both expose ``forward(hidden)->hidden`` and the same key prefix. - self.mlp = Qwen3_5MoE(config, layer_id) if config.moe_enabled else Qwen3_5DenseMLP(config) + self.mlp = Qwen3_5MoE(config, layer_id) if config.moe_enabled else Qwen3_5DenseMLP(config, layer_id) self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -100,6 +100,16 @@ def __init__(self, config: ModelConfig): self.lm_head = Nvfp4LMHead( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size ) + elif getattr(config, "lm_head_quant", "none") == "fp8_pertensor": + # unsloth's mixed export stores lm_head as weight-only FP8 (per-row scale): + # keep it native W8A16 -- the bf16 dequant of this ~2.5 GB matrix would both + # double the memory and be the single largest load-time materialization. + from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorLinear + + assert not config.tie_word_embeddings, "FP8 lm_head assumes untied embeddings" + self.lm_head = Fp8PerTensorLinear( + in_features=config.hidden_size, out_features=config.vocab_size + ) else: self.lm_head = ParallelLMHead( num_embeddings=config.vocab_size, diff --git a/python/freetoken/models/qwen3_5_moe/moe.py b/python/freetoken/models/qwen3_5_moe/moe.py index fc0bb7c2..cb1bb1ce 100644 --- a/python/freetoken/models/qwen3_5_moe/moe.py +++ b/python/freetoken/models/qwen3_5_moe/moe.py @@ -21,13 +21,42 @@ class _SharedExpert(BaseOP): """Always-present shared SwiGLU expert of width ``shared_expert_intermediate_size``.""" - def __init__(self, config: ModelConfig, hidden_size: int, intermediate_size: int): + def __init__(self, config: ModelConfig, hidden_size: int, intermediate_size: int, + layer_id: int | None = None): + # Mixed compressed-tensors exports store some dense-MLP layers in a different + # precision than dense_quant (unsloth: most layers packed NVFP4, a few FP8): + # the per-layer override (dense_mlp_storage) wins when it covers this layer. + storage = "bf16" + overrides = getattr(config, "dense_mlp_storage", None) + if overrides is not None and layer_id is not None and layer_id in overrides: + storage = overrides[layer_id] + elif getattr(config, "dense_quant", "none") == "nvfp4": + storage = "nvfp4" + # storage == "fp8": mixed exports (unsloth) store this dense-MLP layer weight-only + # FP8 (per-row scale) -- keep it native W8A16, the same classes the fp8-split GDN + # uses, instead of the bf16 dequant (halves this layer's memory). + if storage == "fp8": + from freetoken.kernel.triton.fp8_pertensor_linear import ( + Fp8PerTensorColMerged, + Fp8PerTensorLinear, + ) + + self.gate_up_proj = Fp8PerTensorColMerged( + hidden_size, [intermediate_size, intermediate_size], has_bias=False + ) + # down_proj: input = intermediate, output = hidden (in_features, out_features) + # -> weight [hidden, intermediate], matching the other branches + # (Nvfp4DenseLinear / LinearRowParallel both take (intermediate, hidden)). + self.down_proj = Fp8PerTensorLinear( + intermediate_size, hidden_size, has_bias=False + ) + return if getattr(config, "expert_quant", "none") == "fp8_block": self.gate_up_proj = Fp8BlockColMerged( hidden_size, [intermediate_size, intermediate_size], has_bias=False ) self.down_proj = Fp8BlockLinear(intermediate_size, hidden_size, has_bias=False) - elif getattr(config, "dense_quant", "none") == "nvfp4": + elif storage == "nvfp4": # NVFP4 checkpoint: keep the shared expert's NVFP4 weights native (W4A16). from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear @@ -48,12 +77,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Qwen3_5DenseMLP(_SharedExpert): """Dense (non-MoE) SwiGLU MLP for dense Qwen3.x checkpoints (e.g. 27B): ``gate_up_proj`` (fused gate|up) + ``down_proj`` at full ``intermediate_size``. Same structure (and quant - dispatch) as the shared expert -- NVFP4 (W4A16) when ``dense_quant=="nvfp4"``, else bf16 -- - so it reuses ``_SharedExpert`` directly and keeps the state-dict keys flat + dispatch) as the shared expert -- NVFP4 (W4A16) when ``dense_quant=="nvfp4"`` (or the + per-layer ``dense_mlp_storage`` override says so), else bf16 -- so it reuses + ``_SharedExpert`` directly and keeps the state-dict keys flat (``...layers.N.mlp.{gate_up_proj,down_proj}``).""" - def __init__(self, config: ModelConfig): - super().__init__(config, config.hidden_size, config.intermediate_size) + def __init__(self, config: ModelConfig, layer_id: int | None = None): + super().__init__(config, config.hidden_size, config.intermediate_size, layer_id) class Qwen3_5MoE(BaseOP): @@ -73,7 +103,7 @@ def __init__(self, config: ModelConfig, layer_id: int | None = None): ) self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) self.shared_expert = _SharedExpert( - config, config.hidden_size, config.shared_expert_intermediate_size + config, config.hidden_size, config.shared_expert_intermediate_size, layer_id ) self.shared_expert_gate = LinearReplicated(config.hidden_size, 1, has_bias=False) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index b3414089..b569420d 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -87,6 +87,33 @@ def _dequant_fp8_weight(weight: torch.Tensor, weight_scale: torch.Tensor) -> tor return weight.to(torch.bfloat16) * weight_scale.to(torch.bfloat16) +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + + +def _fp8_weight_to_bf16(tensor: torch.Tensor, scale: torch.Tensor | None) -> torch.Tensor: + """Normalize a dense weight stored as fp8 to bf16 before the bf16 fusion/passthrough. + + Mixed NVFP4 exports (e.g. unsloth/Qwen3.8-27B-NVFP4) keep some dense projections + (GDN ``in_proj_*``) as weight-only FP8 instead of bf16; without this, the raw e4m3 + tensor reaches ``ct_bf16_fuse`` and ``torch.cat`` dies on fp8/bf16 promotion. + A per-tensor ``.weight_scale`` makes this a W8A16 dequant (:func:`_dequant_fp8_weight`); + a per-row ``[O, 1]`` scale (unsloth's mixed exports) or block ``[O, IN//g]`` scale is + broadcast over its group; without a scale sibling the fp8 values are the (rounded) + weights themselves, so the cast is exact.""" + if scale is not None: + if scale.numel() == 1: + return _dequant_fp8_weight(tensor, scale) + # Per-row [O, 1] or block [O, IN//g] scale. fp8 values and the (bf16-stored) + # scale are exact in bf16, so a bf16 multiply rounds once -- the same result the + # fp32 path produced before its final cast -- without materializing an fp32 + # copy (a per-row lm_head scale would otherwise transiently double the memory). + if scale.shape[-1] == 1: # per-output-row: pure broadcast, no repeat materialized + return tensor.to(torch.bfloat16) * scale.to(torch.bfloat16) + group = tensor.shape[-1] // scale.shape[-1] + return tensor.to(torch.bfloat16) * scale.to(torch.bfloat16).repeat_interleave(group, dim=-1) + return tensor.to(torch.bfloat16) + + def _dequant_nvfp4_weight( weight: torch.Tensor, weight_scale: torch.Tensor, weight_scale_2: torch.Tensor ) -> torch.Tensor: @@ -180,12 +207,15 @@ def iter_weights( 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. + # Dense compressed-tensors (e.g. Qwen3.6-27B NVFP4; unsloth's mixed NVFP4+FP8 + # export): attn (q/k/v/o, GDN out_proj) + dense MLP W4A16 native where the + # checkpoint stores them packed; FP8 attention/GDN parts stay native W8A16 when + # the model built the fp8-split GDN (attn_quant=="fp8_pertensor"); 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", + attn_fp8=config.attn_quant == "fp8_pertensor", ) return if config.expert_quant == "fp8_block": @@ -311,13 +341,21 @@ def iter_weights( } -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 _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: + """Scalar or per-output-row scale -> per-output-row fp32 vector ``[rows]``. + + A scalar (modelopt's calibrated per-tensor scale) is an exact broadcast; a per-row + ``[rows]``/``[rows, 1]`` scale (unsloth's mixed exports) is used verbatim.""" + scale = scale.to(torch.float32) + if scale.numel() == 1: + return scale.reshape(1).expand(rows) + return scale.reshape(rows) def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor, - act_scale: torch.Tensor | None, buf: dict): + act_scale: torch.Tensor | None, buf: dict, + fuse_map: dict[str, tuple[str, ...]] | None = None, + ) -> list[tuple[str, torch.Tensor]] | list | None: """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. @@ -326,7 +364,7 @@ def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor, 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 fused_suffix, parts in (fuse_map if fuse_map is not None else _PT_FP8_FUSE).items(): for idx, part in enumerate(parts): if base.endswith(part): key = base[: -len(part)] + fused_suffix @@ -564,16 +602,21 @@ def _ct_nvfp4_fuse(base: str, parts_tuple: tuple, buf: dict): def _iter_weights_compressed_tensors( model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool, - nvfp4: bool, + nvfp4: bool, attn_fp8: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: - """Dense pass for a compressed-tensors NVFP4 checkpoint (e.g. Qwen3.6-27B). + """Dense pass for a compressed-tensors checkpoint (dense Qwen3.x, e.g. Qwen3.6-27B + NVFP4, unsloth's mixed NVFP4+FP8 export). 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.""" + With ``attn_fp8`` (unsloth's mixed layout) the attention/GDN projections are kept native + FP8 (W8A16) instead: q/k/v -> ``qkv_proj`` and in_proj_qkv/z -> ``in_proj_qkvz`` (fp8 + + per-row fp32 scale), o_proj / GDN out_proj singletons, in_proj_b/a -> bf16 ``in_proj_ba`` + (matching the model's fp8-split GDN); without it 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: @@ -582,11 +625,30 @@ def _iter_weights_compressed_tensors( tp_info = get_tp_info() nvfp4_buf: dict[str, dict[int, tuple]] = {} bf16_buf: dict[str, dict[int, torch.Tensor]] = {} + fp8_buf: dict[str, dict[int, tuple]] = {} + # The GDN in_proj fusion follows the model layout: the fp8-split GDN (unsloth's mixed + # layout) builds in_proj_qkvz (fp8) + in_proj_ba (bf16); the plain layout fuses all + # four parts into one bf16 in_proj. The mixed layout also stores some dense-MLP + # layers fp8 (native W8A16 gate_up fusion) that the modelopt pass never sees, so the + # fp8 fusion map is extended locally rather than mutating the shared one. + in_proj_fuse = _PT_BF16_FUSE if attn_fp8 else _CT_BF16_FUSE + fp8_fuse = ( + {**_PT_FP8_FUSE, ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj")} + if attn_fp8 + else _PT_FP8_FUSE + ) + # Native-fp8 singletons of the mixed layout (o_proj / GDN out_proj / dense-MLP + # down_proj / lm_head): fp8 weight + per-row fp32 scale, no dequant. + fp8_singletons = ( + (".self_attn.o_proj", ".linear_attn.out_proj", ".mlp.down_proj", "lm_head") + if attn_fp8 + else (".self_attn.o_proj", ".linear_attn.out_proj") + ) 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) + emit = _ct_bf16_fuse(base, tensor, bf16_buf, in_proj_fuse) if emit is not None: yield from emit return @@ -645,7 +707,47 @@ def _emit_bf16_weight(name: str, tensor: torch.Tensor): continue if name.endswith(".weight"): - yield from _emit_bf16_weight(name, reader.get_tensor(raw_name)) + tensor = reader.get_tensor(raw_name) + base = name[: -len(".weight")] + if tensor.dtype in _FP8_DTYPES: + # Mixed exports (unsloth) store some dense projections as weight-only + # FP8 rather than bf16. In the fp8 attention layout everything + # scaled stays native FP8 (W8A16): the fused q/k/v -> qkv_proj, + # in_proj_qkv/z -> in_proj_qkvz, dense-MLP gate/up -> gate_up_proj + # (each part's per-row scale concatenated), the o_proj / GDN + # out_proj / dense-MLP down_proj / lm_head singletons. Unscaled fp8 + # dequantizes to bf16 -- the bf16 fusion's torch.cat dies on raw + # fp8/bf16 promotion. The scale sibling is skipped globally by + # this pass, so this is its only consumer. + raw_base = raw_name[: -len(".weight")] + scale_name = raw_base + ".weight_scale" + has_scale = reader.has(scale_name) + scale = reader.get_tensor(scale_name) if has_scale else None + if attn_fp8 and has_scale: + emit = _pt_fp8_fuse(base, tensor, scale, None, fp8_buf, fp8_fuse) + if emit is not None: + # The fp8 fusion owns this part (buffered `[]` or just + # completed): it must never also reach the bf16 side, or + # the group completes on the fp8 side while the bf16 side + # is left short one part (incomplete bf16 fusion assert). + yield from emit + continue + if base.endswith(fp8_singletons): + yield base + ".weight", tensor + yield ( + base + ".weight_scale", + _per_row_scale(scale, tensor.shape[0]).contiguous(), + ) + continue + tensor = _fp8_weight_to_bf16(tensor, scale) + # bf16 parts (and dequantized fp8) feed the bf16 fusions: q/k/v -> + # qkv_proj, gate/up -> gate_up_proj (NVFP4 layout map), plus the + # layout's in_proj form; singletons / norms pass through. + emit = _ct_bf16_fuse(base, tensor, bf16_buf, _CT_NVFP4_FUSE) + if emit is None: + yield from _emit_bf16_weight(name, tensor) + else: + yield from emit continue # A_log / dt_bias (kept fp32 by the model; the load downcast exempts them). @@ -655,6 +757,7 @@ def _emit_bf16_weight(name: str, tensor: torch.Tensor): assert not nvfp4_buf, f"Incomplete NVFP4 fusions: {list(nvfp4_buf.keys())}" assert not bf16_buf, f"Incomplete bf16 fusions: {list(bf16_buf.keys())}" + assert not fp8_buf, f"Incomplete fp8 fusions: {list(fp8_buf.keys())}" def iter_weights_parallel( diff --git a/tests/models/test_qwen3_5_moe_config.py b/tests/models/test_qwen3_5_moe_config.py new file mode 100644 index 00000000..b8f4d50a --- /dev/null +++ b/tests/models/test_qwen3_5_moe_config.py @@ -0,0 +1,167 @@ +"""``_compressed_linear_storage`` (weight_map sniffing for mixed compressed-tensors +exports, e.g. unsloth/Qwen3.8-27B-NVFP4) and the per-layer dense-MLP storage gate. + +The dense pass keeps each linear in the storage the checkpoint actually uses: native +W4A16 for packed layers, native W8A16 for fp8 attention linears, bf16 dequant for the +rest. When the index is unavailable the fallback must preserve the official +dense-NVFP4 assumption (all packed, no overrides).""" + +import json +import os +import tempfile +import unittest + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.models.qwen3_5_moe.config import _compressed_linear_storage + +if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +class _FakeHFConfig: + def __init__(self, name_or_path=None): + self.name_or_path = name_or_path + + +def _write_index(tmpdir, keys): + with open(os.path.join(tmpdir, "model.safetensors.index.json"), "w", encoding="utf-8") as f: + json.dump({"metadata": {}, "weight_map": {k: "shard.bin" for k in keys}}, f) + + +class _CompressedLinearStorageTest(unittest.TestCase): + def test_mixed_unsloth_layout(self): + # unsloth/Qwen3.8-27B-NVFP4: attention linears FP8 (weight + scale), most mlp + # NVFP4 (weight_packed), layers 56-63 mlp FP8 (weight + scale), lm_head FP8. + keys = [] + for layer in range(64): + keys += [ + f"model.language_model.layers.{layer}.self_attn.q_proj.weight", + f"model.language_model.layers.{layer}.self_attn.q_proj.weight_scale", + f"model.language_model.layers.{layer}.linear_attn.out_proj.weight", + f"model.language_model.layers.{layer}.linear_attn.out_proj.weight_scale", + ] + if layer < 56: + keys.append(f"model.language_model.layers.{layer}.mlp.gate_proj.weight_packed") + else: + keys.append(f"model.language_model.layers.{layer}.mlp.gate_proj.weight") + keys.append(f"model.language_model.layers.{layer}.mlp.gate_proj.weight_scale") + keys.append("lm_head.weight") + keys.append("lm_head.weight_scale") + with tempfile.TemporaryDirectory() as tmp: + _write_index(tmp, keys) + attn, dense_fallback, overrides, lmhead_fp8 = _compressed_linear_storage( + _FakeHFConfig(tmp) + ) + self.assertEqual(attn, "fp8") + self.assertEqual(dense_fallback, "nvfp4") + self.assertEqual(overrides, {l: "fp8" for l in range(56, 64)}) + self.assertTrue(lmhead_fp8) + + def test_official_dense_nvfp4_kept_native(self): + keys = [ + "model.layers.0.self_attn.q_proj.weight_packed", + "model.layers.0.self_attn.q_proj.weight_scale", + "model.layers.0.linear_attn.out_proj.weight_packed", + "model.layers.0.mlp.gate_proj.weight_packed", + "model.layers.0.mlp.up_proj.weight_packed", + "model.layers.0.mlp.down_proj.weight_packed", + "lm_head.weight", + ] + with tempfile.TemporaryDirectory() as tmp: + _write_index(tmp, keys) + self.assertEqual( + _compressed_linear_storage(_FakeHFConfig(tmp)), + ("nvfp4", "nvfp4", None, False), + ) + + def test_routed_expert_moe_naming_keeps_native(self): + # MoE checkpoint: dense linears live under mlp.shared_expert / experts -- no + # mlp.{gate,up,down}_proj keys at all -> no overrides, native assumption. + keys = [ + "model.layers.0.self_attn.q_proj.weight_packed", + "model.layers.0.mlp.shared_expert.gate_proj.weight_packed", + "model.layers.0.mlp.experts.0.gate_proj.weight_packed", + ] + with tempfile.TemporaryDirectory() as tmp: + _write_index(tmp, keys) + self.assertEqual( + _compressed_linear_storage(_FakeHFConfig(tmp)), + ("nvfp4", "nvfp4", None, False), + ) + + def test_index_unavailable_falls_back_to_native(self): + # Hub id before download (name_or_path not a local dir) and missing index file. + self.assertEqual( + _compressed_linear_storage(_FakeHFConfig("unsloth/Qwen3.8-27B-NVFP4")), + ("nvfp4", "nvfp4", None, False), + ) + with tempfile.TemporaryDirectory() as tmp: + self.assertEqual( + _compressed_linear_storage(_FakeHFConfig(tmp)), + ("nvfp4", "nvfp4", None, False), + ) + + +class _FakeConfig: + def __init__(self, **kw): + self.expert_quant = kw.get("expert_quant", "none") + self.dense_quant = kw.get("dense_quant", "none") + self.dense_mlp_storage = kw.get("dense_mlp_storage", None) + + +class _SharedExpertStorageTest(unittest.TestCase): + """The per-layer dense-MLP storage gate: the dense_mlp_storage override wins over + the dense_quant flag for the layers it covers.""" + + def _expert(self, layer_id=None, **kw): + from freetoken.models.qwen3_5_moe.moe import _SharedExpert + + # in_features must be %16 for the NVFP4 dense kernels + return _SharedExpert(_FakeConfig(**kw), 16, 32, layer_id) + + def test_per_layer_bf16_override_wins(self): + from freetoken.layers import LinearColParallelMerged + + e = self._expert( + layer_id=0, dense_quant="nvfp4", dense_mlp_storage={0: "bf16"} + ) + self.assertIsInstance(e.gate_up_proj, LinearColParallelMerged) + + def test_per_layer_fp8_override_stays_native(self): + from freetoken.kernel.triton.fp8_pertensor_linear import ( + Fp8PerTensorColMerged, + Fp8PerTensorLinear, + ) + + e = self._expert( + layer_id=60, dense_quant="nvfp4", dense_mlp_storage={60: "fp8"} + ) + self.assertIsInstance(e.gate_up_proj, Fp8PerTensorColMerged) + self.assertIsInstance(e.down_proj, Fp8PerTensorLinear) + # Orientation: gate/up read hidden -> write intermediate; down reads intermediate + # -> writes hidden (weight rows = output dim). + self.assertEqual(tuple(e.gate_up_proj.weight.shape), (64, 16)) + self.assertEqual(tuple(e.down_proj.weight.shape), (16, 32)) + + def test_packed_layer_stays_native(self): + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged + + e = self._expert( + layer_id=0, dense_quant="nvfp4", dense_mlp_storage={1: "bf16"} + ) + self.assertIsInstance(e.gate_up_proj, Nvfp4DenseColMerged) + + def test_flag_fallback_without_map(self): + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged + from freetoken.layers import LinearColParallelMerged + + self.assertIsInstance( + self._expert(dense_quant="nvfp4").gate_up_proj, Nvfp4DenseColMerged + ) + self.assertIsInstance( + self._expert().gate_up_proj, LinearColParallelMerged + ) + + +if __name__ == "__main__": + unittest.main() 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 00000000..ca10dd2b --- /dev/null +++ b/tests/models/test_qwen3_5_moe_weight.py @@ -0,0 +1,211 @@ +"""qwen3_5_moe dense-pass normalization for mixed-precision NVFP4 exports. + +unsloth/Qwen3.8-27B-NVFP4 keeps some dense projections (GDN ``in_proj_*``) as +weight-only FP8 instead of bf16. The dense pass used to feed the raw e4m3 tensor +straight into ``ct_bf16_fuse``, whose ``torch.cat`` dies on fp8/bf16 promotion: + + RuntimeError: Promotion for Float8 Types is not supported, attempted to promote + Float8_e4m3fn and BFloat16 + +The pass now normalizes fp8 ``.weight`` tensors to bf16 first (W8A16 dequant with a +per-tensor scale, broadcast with a block scale, exact cast without a scale). +""" + +import torch + +from freetoken.models.loader import ct_bf16_fuse +from freetoken.models.qwen3_5_moe.weight import ( + _CT_BF16_FUSE, + _CT_NVFP4_FUSE, + _FP8_DTYPES, + _fp8_weight_to_bf16, +) + + +def test_fp8_with_per_tensor_scale_dequants(): + w = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32).to(torch.float8_e4m3fn) + scale = torch.tensor(0.25) + out = _fp8_weight_to_bf16(w, scale) + assert out.dtype == torch.bfloat16 + assert torch.equal(out, w.to(torch.bfloat16) * scale.to(torch.bfloat16)) + + +def test_fp8_with_scalar_shape1_scale_dequants(): + w = torch.tensor([[1.0, 2.0]], dtype=torch.float32).to(torch.float8_e4m3fn) + out = _fp8_weight_to_bf16(w, torch.tensor([2.0])) + assert out.dtype == torch.bfloat16 + assert torch.equal(out, w.to(torch.bfloat16) * torch.tensor([2.0]).to(torch.bfloat16)) + + +def test_fp8_without_scale_is_exact_cast(): + w = torch.tensor([1.0, 2.0, 3.5], dtype=torch.float32).to(torch.float8_e4m3fn) + out = _fp8_weight_to_bf16(w, None) + assert out.dtype == torch.bfloat16 + assert torch.equal(out, w.to(torch.bfloat16)) + + +def test_fp8_with_block_scale_broadcasts(): + # [O, IN] weight, block scale [O, IN//2] -> group size 2. fp8 values and the scale + # are exact in bf16, so the bf16 multiply rounds once (same result as the fp32 + # product before its final cast) without a materialized fp32 copy. + w = torch.arange(8, dtype=torch.float32).reshape(2, 4).to(torch.float8_e4m3fn) + s = torch.tensor([[1.0, 2.0], [0.5, 4.0]], dtype=torch.bfloat16) + out = _fp8_weight_to_bf16(w, s) + assert out.shape == (2, 4) + assert out.dtype == torch.bfloat16 + expected = w.to(torch.bfloat16) * s.to(torch.bfloat16).repeat_interleave(2, dim=1) + assert torch.equal(out, expected) + + +def test_fp8_with_per_row_scale_broadcasts(): + # unsloth's mixed exports store a per-row [O, 1] scale: pure broadcast multiply + # (no repeat materialized -- a per-row lm_head scale would otherwise transiently + # double the memory of the dequant). + w = torch.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], dtype=torch.float32).to( + torch.float8_e4m3fn + ) + s = torch.tensor([[0.5], [2.0]], dtype=torch.bfloat16) + out = _fp8_weight_to_bf16(w, s) + assert out.shape == (2, 4) + assert out.dtype == torch.bfloat16 + expected = w.to(torch.bfloat16) * s.to(torch.bfloat16) + assert torch.equal(out, expected) + + +def test_bf16_tensor_untouched_by_dtype_gate(): + """The call site only normalizes fp8 dtypes; bf16 parts flow through as-is.""" + w = torch.randn(4, 4, dtype=torch.bfloat16) + assert w.dtype not in _FP8_DTYPES + + +def test_mixed_group_bf16_fusion_repro(): + """Regression: one fp8 part + bf16 parts in the GDN ``in_proj`` group. Without the + normalization the group's torch.cat raises 'Promotion for Float8 Types is not + supported'; with it, the fused ``in_proj.weight`` is bf16 and output-dim-concatenated + in the canonical (qkv, z, b, a) order.""" + O, IN = 8, 4 + # All parts share the input dim (last dim) and differ only in output rows, like a + # real in_proj group. + parts = ( + (".linear_attn.in_proj_qkv", torch.randn(O, IN).to(torch.float8_e4m3fn)), + (".linear_attn.in_proj_z", torch.randn(O, IN).to(torch.bfloat16)), + (".linear_attn.in_proj_b", torch.randn(O, IN).to(torch.bfloat16)), + (".linear_attn.in_proj_a", torch.randn(O, IN).to(torch.bfloat16)), + ) + base = "model.layers.0" + buf: dict = {} + completed: list[tuple[str, torch.Tensor]] = [] + for part, t in parts: + # The dense pass normalizes fp8 parts before fusing (the fix under test). + if t.dtype in _FP8_DTYPES: + t = _fp8_weight_to_bf16(t, None) + emit = ct_bf16_fuse(base + part, t, buf, _CT_BF16_FUSE) + if emit: + completed.extend(emit) + assert len(completed) == 1 + (key, fused), = completed + assert key == base + ".linear_attn.in_proj.weight" + assert fused.dtype == torch.bfloat16 + assert fused.shape == (O * 4, IN) + assert not buf + + +def test_fp8_qkv_parts_fuse_to_qkv_proj(): + """unsloth's self_attn q/k/v are weight-only FP8: after normalization they must fuse + into qkv_proj through the same bf16 fusion map the NVFP4 dequant branch uses. + The model builds a bf16 qkv linear for mixed layouts (dense pass emits qkv_proj).""" + O, IN = 4, 4 + base = "model.layers.0.self_attn" + buf: dict = {} + completed: list[tuple[str, torch.Tensor]] = [] + for part in (".q_proj", ".k_proj", ".v_proj"): + t = _fp8_weight_to_bf16(torch.randn(O, IN).to(torch.float8_e4m3fn), None) + emit = ct_bf16_fuse(base + part, t, buf, _CT_NVFP4_FUSE) + if emit: + completed.extend(emit) + assert len(completed) == 1 + (key, fused), = completed + assert key == base + ".qkv_proj.weight" + assert fused.dtype == torch.bfloat16 + assert fused.shape == (O * 3, IN) + assert not buf + + +def test_fp8_gate_up_parts_fuse_to_gate_up_proj(): + """unsloth stores some dense-MLP layers' gate/up as weight-only FP8: both must fuse + into gate_up_proj (the model builds a bf16 gate_up linear for mixed layouts).""" + O, IN = 6, 4 + base = "model.layers.0.mlp" + buf: dict = {} + completed: list[tuple[str, torch.Tensor]] = [] + for part in (".gate_proj", ".up_proj"): + t = _fp8_weight_to_bf16(torch.randn(O, IN).to(torch.float8_e4m3fn), None) + emit = ct_bf16_fuse(base + part, t, buf, _CT_NVFP4_FUSE) + if emit: + completed.extend(emit) + assert len(completed) == 1 + (key, fused), = completed + assert key == base + ".gate_up_proj.weight" + assert fused.shape == (O * 2, IN) + assert not buf + + +def test_fp8_o_proj_not_fused_stays_standalone(): + """o_proj / out_proj are singletons in every layout: the qkv/gate_up fusion map must + not swallow them (they pass through as bare .weight).""" + for part in (".o_proj", ".out_proj"): + t = _fp8_weight_to_bf16(torch.randn(4, 4).to(torch.float8_e4m3fn), None) + buf: dict = {} + assert ct_bf16_fuse("model.layers.0" + part, t, buf, _CT_NVFP4_FUSE) is None + assert not buf + + +def test_per_row_scale_used_verbatim(): + from freetoken.models.qwen3_5_moe.weight import _per_row_scale + + s = torch.tensor([[0.25], [0.5]], dtype=torch.bfloat16) + out = _per_row_scale(s, 2) + assert out.dtype == torch.float32 + assert out.shape == (2,) + assert torch.equal(out, torch.tensor([0.25, 0.5], dtype=torch.float32)) + + +def test_scalar_scale_still_broadcasts(): + from freetoken.models.qwen3_5_moe.weight import _per_row_scale + + out = _per_row_scale(torch.tensor([2.0]), 3) + assert torch.equal(out, torch.full((3,), 2.0, dtype=torch.float32)) + + +def test_fp8_native_qkv_fusion_per_row_scales(): + """unsloth's attention: fp8 q/k/v with per-row [O, 1] scales stay native FP8 (W8A16) + -- the fusion concatenates the fp8 weights and the per-row fp32 scales, no dequant.""" + from freetoken.models.qwen3_5_moe.weight import _pt_fp8_fuse + + O, IN = 4, 8 + base = "model.layers.3.self_attn" + wq = torch.randn(O, IN).to(torch.float8_e4m3fn) + wk = torch.randn(O, IN).to(torch.float8_e4m3fn) + wv = torch.randn(O, IN).to(torch.float8_e4m3fn) + sq = torch.tensor([[0.25], [0.5], [1.0], [0.125]], dtype=torch.bfloat16) + sk = torch.tensor([[0.5], [0.25], [2.0], [1.0]], dtype=torch.bfloat16) + sv = torch.tensor([[1.0], [1.0], [1.0], [1.0]], dtype=torch.bfloat16) + buf: dict = {} + completed: list[tuple[str, torch.Tensor]] = [] + for part, w, s in (".q_proj", wq, sq), (".k_proj", wk, sk), (".v_proj", wv, sv): + emit = _pt_fp8_fuse(base + part, w, s, None, buf) + if emit: + completed.extend(emit) + assert len(completed) == 2 # weight + weight_scale (no input_scale: acts are None) + (key_w, w), (key_s, s) = completed + assert key_w == base + ".qkv_proj.weight" + assert w.dtype == torch.float8_e4m3fn + assert w.shape == (O * 3, IN) + assert key_s == base + ".qkv_proj.weight_scale" + assert s.dtype == torch.float32 + assert s.shape == (O * 3,) + expected = torch.cat( + [sq.reshape(-1).float(), sk.reshape(-1).float(), sv.reshape(-1).float()] + ) + assert torch.equal(s, expected) + assert not buf From ed9ddfe61592de54681d2e568c57f6720c9decb4 Mon Sep 17 00:00:00 2001 From: Chris Qian Date: Wed, 26 Aug 2026 12:25:44 +0800 Subject: [PATCH 2/3] docs: add unsloth Qwen3.8 dense NVFP4 (mixed-precision export) to supported models unsloth's per-module mixed-precision dense exports (NVFP4 MLP + FP8 attention/GDN/lm_head + bf16 residual parts) load natively end-to-end; list the known-good checkpoint and document the layout. --- docs/models.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/models.md b/docs/models.md index e4850a12..80b5a782 100644 --- a/docs/models.md +++ b/docs/models.md @@ -11,6 +11,7 @@ for them; other checkpoints of the same architectures work too. | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) | | Qwen3.6 / Qwen3.5 MoE | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-35B-A3B-FP8)), [nvidia/Qwen3.6-35B-A3B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-35B-A3B-NVFP4), [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8)) | | Qwen3.6 dense | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-27B-FP8)), [nvidia/Qwen3.6-27B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-27B-NVFP4) | +| Qwen3.8 dense | [unsloth/Qwen3.8-27B-NVFP4](https://huggingface.co/unsloth/Qwen3.8-27B-NVFP4) (mixed NVFP4+FP8 export) | | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | | Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. | @@ -38,3 +39,6 @@ for them; other checkpoints of the same architectures work too. - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the authoritative model args are read from there. - Multimodal checkpoints are served text-only. +- unsloth's dense Qwen3.x NVFP4 exports are mixed-precision per module (NVFP4 MLP + layers, FP8 attention/GDN/lm_head, bf16 residual parts); each linear loads in the + storage the checkpoint actually uses. From e8bca65a717c2b07aa22567cd487c67e624c8563 Mon Sep 17 00:00:00 2001 From: Chris Qian Date: Wed, 26 Aug 2026 13:38:02 +0800 Subject: [PATCH 3/3] qwen3_5_moe: tidy the mixed-precision dense pass No behavior change; the mixed-NVFP4+FP8 load path (unsloth/Qwen3.8-27B-NVFP4) stays byte-identical: - _pt_fp8_fuse: rename the 'scalar' param to 'scale' (it now accepts modelopt scalars AND unsloth's per-row [O, 1] scales) and fix the sloppy return annotation (bare 'list' -> list[tuple[str, torch.Tensor]] | None) - _per_row_scale: fail loud with a clear message if a non-scalar scale has the wrong element count instead of a cryptic reshape error - document the unscaled-fp8 fallthrough assumption in the .weight handler (an unscaled fp8 q/k/v or in_proj_qkv/z in the fp8-split layout would fail at load with a missing key, not silently dequant -- no real export has this) - weight tests: refresh the module docstring (native W8A16 + dequant, not dequant-only) and pin the local fp8-fusion-map extension (dense-MLP gate/up -> gate_up_proj native, per-row fp32 scales) with a unit test Verified: config tests 8/8, weight tests 14/14, parse_config on the real checkpoint, full-pass key/shape/dtype check (0 stray / 0 missing / 0 never-yielded) all green on the remote (RTX 5090 D, editable install). Co-Authored-By: GooeyPi --- python/freetoken/models/qwen3_5_moe/weight.py | 20 +++++++-- tests/models/test_qwen3_5_moe_weight.py | 41 ++++++++++++++++--- 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index b569420d..50457c93 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -349,17 +349,24 @@ def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: scale = scale.to(torch.float32) if scale.numel() == 1: return scale.reshape(1).expand(rows) + if scale.numel() != rows: + raise ValueError( + f"per-row scale must have {rows} elements, got {scale.numel()} (shape {tuple(scale.shape)})" + ) return scale.reshape(rows) -def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor, +def _pt_fp8_fuse(base: str, weight: torch.Tensor, scale: torch.Tensor, act_scale: torch.Tensor | None, buf: dict, fuse_map: dict[str, tuple[str, ...]] | None = None, - ) -> list[tuple[str, torch.Tensor]] | list | None: - """Buffer an fp8 fusion part ``(weight, scalar, act_scale)``; once all parts arrive emit + ) -> list[tuple[str, torch.Tensor]] | None: + """Buffer an fp8 fusion part ``(weight, scale, 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. + ``scale`` is the part's ``.weight_scale``: a calibrated scalar (modelopt) or a per-row + ``[O, 1]``/``[O]`` vector (unsloth's mixed exports) -- normalized per row downstream. + 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 @@ -369,7 +376,7 @@ def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor, if base.endswith(part): key = base[: -len(part)] + fused_suffix slots = buf.setdefault(key, {}) - slots[idx] = (weight, scalar, act_scale) + slots[idx] = (weight, scale, act_scale) if len(slots) < len(parts): return [] del buf[key] @@ -739,6 +746,11 @@ def _emit_bf16_weight(name: str, tensor: torch.Tensor): _per_row_scale(scale, tensor.shape[0]).contiguous(), ) continue + # Unscaled fp8 dequantizes to bf16. Note: an *unscaled* fp8 + # q/k/v or in_proj_qkv/z in the fp8-split layout (no real export + # has this: unsloth scales every fp8 part) would dequant into the + # bf16 group instead of the model's native-fp8 linears and fail at + # load with a missing key -- loud, not a silent dequant. tensor = _fp8_weight_to_bf16(tensor, scale) # bf16 parts (and dequantized fp8) feed the bf16 fusions: q/k/v -> # qkv_proj, gate/up -> gate_up_proj (NVFP4 layout map), plus the diff --git a/tests/models/test_qwen3_5_moe_weight.py b/tests/models/test_qwen3_5_moe_weight.py index ca10dd2b..859eb59a 100644 --- a/tests/models/test_qwen3_5_moe_weight.py +++ b/tests/models/test_qwen3_5_moe_weight.py @@ -1,14 +1,17 @@ """qwen3_5_moe dense-pass normalization for mixed-precision NVFP4 exports. -unsloth/Qwen3.8-27B-NVFP4 keeps some dense projections (GDN ``in_proj_*``) as -weight-only FP8 instead of bf16. The dense pass used to feed the raw e4m3 tensor -straight into ``ct_bf16_fuse``, whose ``torch.cat`` dies on fp8/bf16 promotion: +unsloth/Qwen3.8-27B-NVFP4 keeps some dense projections as weight-only FP8 instead of +bf16. The dense pass used to feed the raw e4m3 tensor straight into ``ct_bf16_fuse``, +whose ``torch.cat`` dies on fp8/bf16 promotion: RuntimeError: Promotion for Float8 Types is not supported, attempted to promote Float8_e4m3fn and BFloat16 -The pass now normalizes fp8 ``.weight`` tensors to bf16 first (W8A16 dequant with a -per-tensor scale, broadcast with a block scale, exact cast without a scale). +The pass now keeps scaled fp8 linears native where the model built fp8 linears (W8A16: +fp8 weight + per-row fp32 scale, fused q/k/v -> qkv_proj, in_proj_qkv/z -> in_proj_qkvz, +gate/up -> gate_up_proj, o_proj / GDN out_proj / down_proj / lm_head singletons) and +dequantizes the remainder to bf16 first (W8A16 dequant with a per-tensor scale, +broadcast with a block scale, exact cast without a scale). """ import torch @@ -177,6 +180,34 @@ def test_scalar_scale_still_broadcasts(): assert torch.equal(out, torch.full((3,), 2.0, dtype=torch.float32)) +def test_fp8_mlp_gate_up_fusion_local_map(): + """The mixed layout extends the fp8 fusion map locally (dense-MLP gate/up stay native + W8A16 instead of dequantizing to bf16): fp8 weights + concatenated per-row fp32 + scales, no input_scale (acts are None).""" + from freetoken.models.qwen3_5_moe.weight import _PT_FP8_FUSE, _pt_fp8_fuse + + O, IN = 4, 8 + base = "model.layers.56.mlp" + local_map = {**_PT_FP8_FUSE, ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj")} + w = {p: torch.randn(O, IN).to(torch.float8_e4m3fn) for p in (".gate_proj", ".up_proj")} + s = {p: torch.full((O, 1), 0.5, dtype=torch.bfloat16) for p in w} + buf: dict = {} + completed: list[tuple[str, torch.Tensor]] = [] + for part in (".gate_proj", ".up_proj"): + emit = _pt_fp8_fuse(base + part, w[part], s[part], None, buf, local_map) + if emit: + completed.extend(emit) + assert len(completed) == 2 # weight + weight_scale (no input_scale: acts are None) + (key_w, fused_w), (key_s, fused_s) = completed + assert key_w == base + ".gate_up_proj.weight" + assert fused_w.dtype == torch.float8_e4m3fn + assert fused_w.shape == (O * 2, IN) + assert key_s == base + ".gate_up_proj.weight_scale" + assert fused_s.dtype == torch.float32 + assert fused_s.shape == (O * 2,) + assert not buf + + def test_fp8_native_qkv_fusion_per_row_scales(): """unsloth's attention: fp8 q/k/v with per-row [O, 1] scales stay native FP8 (W8A16) -- the fusion concatenates the fp8 weights and the per-row fp32 scales, no dequant."""