From 66b957e2fc2808db686a37def6aafdfa329877c0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 19:53:36 -0700 Subject: [PATCH 1/6] convert: support compressed-tensors mixed-precision NVFP4 checkpoints Every NVFP4 checkpoint published by Unsloth (unsloth/Qwen3.6-35B-A3B-NVFP4, unsloth/Qwen3.6-35B-A3B-NVFP4-Fast, unsloth/Qwen3.6-27B-NVFP4) uses the compressed-tensors "mixed-precision" format with two config groups: one float-quantized FP8 group covering the attention projections and lm_head, and one nvfp4-pack-quantized group covering the MoE experts. The converter rejected all of them with NotImplementedError: Can't handle multiple config groups for compressed-tensors yet because the nvfp4_compressed_tensors gate required every group to be nvfp4-pack-quantized. Relax both copies of that gate to accept a checkpoint in which any group is NVFP4, and handle the rest of the checkpoint: - _generate_nvfp4_tensors now identifies NVFP4 tensors by dtype and block geometry rather than by scale rank alone. The FP8 group also carries a 2D weight_scale of shape [out, 1], so the existing "scale.ndim < 2" test let FP8 tensors fall into the NVFP4 repacking path. - The nvfp4 branch of dequant_model now dequantizes the leftover FP8 weights the same way the float-quantized branch does, and drops the unused input_scale, k_scale and v_scale sidecars. Previously it did nothing, so those tensors reached the writer still quantized. With this, converting unsloth/Qwen3.6-35B-A3B-NVFP4-Fast with --fp8-as-q8 produces a MOSTLY_NVFP4 GGUF whose 240 expert tensors are GGML_TYPE_NVFP4 and whose attention and lm_head tensors are Q8_0. --- conversion/base.py | 36 +++++++++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 56547ace009f..81b03ecfc283 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -489,7 +489,7 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T quant_format == "nvfp4-pack-quantized" or quant_format == "mixed-precision" and bool(groups) - and all(g.get("format") == "nvfp4-pack-quantized" for g in groups.values() if isinstance(g, dict)) + and any(g.get("format") == "nvfp4-pack-quantized" for g in groups.values() if isinstance(g, dict)) ) if len(groups) > 1 and not nvfp4_compressed_tensors: @@ -538,8 +538,27 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T if (base_name + "_zero_point") in self.model_tensors: tensors_to_remove.append(base_name + "_zero_point") elif nvfp4_compressed_tensors: - # Don't error from compressed-tensors, we'll handle them in _generate_nvfp4_tensors - pass + # NVFP4 tensors were already repacked by _generate_nvfp4_tensors and removed + # from model_tensors. For a "mixed-precision" checkpoint whatever weight_scale + # entries are left belong to the non-NVFP4 config group (FP8 per-channel); + # dequantize them exactly like the float-quantized branch above. + for name in self.model_tensors.keys(): + if name.endswith(".weight_scale"): + weight_name = name.removesuffix("_scale") + if weight_name not in self.model_tensors: + tensors_to_remove.append(name) + continue + w = self.model_tensors[weight_name] + s = self.model_tensors[name] + is_fp8 = False + if self._fp8_as_q8: + is_fp8 = w().dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + self.model_tensors[weight_name] = lambda w=w, s=s: dequant_simple(w(), s(), None) + tensors_to_remove.append(name) + if is_fp8: + self._fp8_dequantized.add(weight_name) + elif name.endswith((".input_scale", ".k_scale", ".v_scale", ".weight_scale_2")): + tensors_to_remove.append(name) else: raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported") elif quant_method == "modelopt": @@ -751,9 +770,16 @@ def _generate_nvfp4_tensors(self): weight = LazyTorchTensor.to_eager(self.model_tensors[name]()) scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]()) - # Skip non-NVFP4 tensors (e.g. FP8 with per-channel 1D scales) + # Skip non-NVFP4 tensors (e.g. FP8 with per-channel 1D scales). + # In a compressed-tensors "mixed-precision" checkpoint the FP8 group also has a + # 2D weight_scale of shape [out, 1], so shape alone is not enough: an NVFP4 + # tensor is nibble-packed uint8 with an E4M3 scale, one per 16 values. if scale.ndim < 2: continue + if weight.dtype != torch.uint8 or scale.dtype != torch.float8_e4m3fn: + continue + if scale.shape[-1] * 16 != weight.shape[-1] * 2: + continue scale2 = LazyTorchTensor.to_eager(self.model_tensors.get(scale2_name, lambda: torch.tensor(1.0))()) input_scale = LazyTorchTensor.to_eager(self.model_tensors.get(input_scale_name, lambda: torch.tensor(1.0))()) @@ -858,7 +884,7 @@ def prepare_tensors(self): quant_format == "nvfp4-pack-quantized" or quant_format == "mixed-precision" and bool(quant_groups) - and all(g.get("format") == "nvfp4-pack-quantized" for g in quant_groups.values() if isinstance(g, dict)) + and any(g.get("format") == "nvfp4-pack-quantized" for g in quant_groups.values() if isinstance(g, dict)) ) if quant_algo != "NVFP4": if nvfp4_compressed_tensors: From 9a03f5c6422859f6ef164c4e8369814b25dd9a88 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 18:37:44 -0700 Subject: [PATCH 2/6] convert: trim comments in the NVFP4 mixed-precision changes --- conversion/base.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 81b03ecfc283..8a6e1a2c8f2a 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -538,10 +538,7 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T if (base_name + "_zero_point") in self.model_tensors: tensors_to_remove.append(base_name + "_zero_point") elif nvfp4_compressed_tensors: - # NVFP4 tensors were already repacked by _generate_nvfp4_tensors and removed - # from model_tensors. For a "mixed-precision" checkpoint whatever weight_scale - # entries are left belong to the non-NVFP4 config group (FP8 per-channel); - # dequantize them exactly like the float-quantized branch above. + # _generate_nvfp4_tensors already removed the NVFP4 tensors, so a leftover weight_scale is the FP8 group. for name in self.model_tensors.keys(): if name.endswith(".weight_scale"): weight_name = name.removesuffix("_scale") @@ -770,10 +767,7 @@ def _generate_nvfp4_tensors(self): weight = LazyTorchTensor.to_eager(self.model_tensors[name]()) scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]()) - # Skip non-NVFP4 tensors (e.g. FP8 with per-channel 1D scales). - # In a compressed-tensors "mixed-precision" checkpoint the FP8 group also has a - # 2D weight_scale of shape [out, 1], so shape alone is not enough: an NVFP4 - # tensor is nibble-packed uint8 with an E4M3 scale, one per 16 values. + # ndim alone is not enough: the FP8 group also has a 2D [out, 1] weight_scale. NVFP4 is nibble-packed uint8, E4M3 scale, one per 16. if scale.ndim < 2: continue if weight.dtype != torch.uint8 or scale.dtype != torch.float8_e4m3fn: From 4a4dd203d320be62a7ab87da8be6e7399e97b9f0 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 18:51:02 -0700 Subject: [PATCH 3/6] convert: refuse NVFP4 mixed-precision paired with a block-quantized group The residual branch dequantized every leftover weight_scale with block_size None, which is right for the per-channel FP8 group this path was written for and wrong for a block group, whose scales are a grid that block_structure has to expand first. Using them directly either raises on broadcasting or, when the dimensions happen to line up, applies the wrong scales silently. Resolve each residual group and reject anything but the supported channel pairing. --- conversion/base.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/conversion/base.py b/conversion/base.py index 8a6e1a2c8f2a..e7925aa473f5 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -539,6 +539,17 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T tensors_to_remove.append(base_name + "_zero_point") elif nvfp4_compressed_tensors: # _generate_nvfp4_tensors already removed the NVFP4 tensors, so a leftover weight_scale is the FP8 group. + # Only per-channel residual groups are handled: a "block" group's scales are a grid needing + # block_structure expansion, and passing them straight to dequant_simple misapplies them. + for group in groups.values(): + if not isinstance(group, dict) or group.get("format") == "nvfp4-pack-quantized": + continue + residual = group.get("weights") or {} + if residual.get("strategy") != "channel" or residual.get("block_structure") is not None: + raise NotImplementedError( + f"compressed-tensors mixed-precision with NVFP4 plus a " + f"{residual.get('strategy')!r} group is not yet supported" + ) for name in self.model_tensors.keys(): if name.endswith(".weight_scale"): weight_name = name.removesuffix("_scale") From fc33fac7ab49f43c9b584f8c7ab6f7a4d420e326 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 21:07:18 -0700 Subject: [PATCH 4/6] convert: tighten the NVFP4 mixed-precision comments Post-convergence pass. Four comment lines to three, and the two over-long ones brought under the line limit. Comments only, no code change. --- conversion/base.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index e7925aa473f5..ffc0e8bd5516 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -538,9 +538,9 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T if (base_name + "_zero_point") in self.model_tensors: tensors_to_remove.append(base_name + "_zero_point") elif nvfp4_compressed_tensors: - # _generate_nvfp4_tensors already removed the NVFP4 tensors, so a leftover weight_scale is the FP8 group. - # Only per-channel residual groups are handled: a "block" group's scales are a grid needing - # block_structure expansion, and passing them straight to dequant_simple misapplies them. + # NVFP4 tensors are gone by here, so a leftover weight_scale is the FP8 group. + # Per-channel residuals only: a block group's scales are a grid needing + # block_structure expansion, which dequant_simple would misapply. for group in groups.values(): if not isinstance(group, dict) or group.get("format") == "nvfp4-pack-quantized": continue @@ -778,7 +778,7 @@ def _generate_nvfp4_tensors(self): weight = LazyTorchTensor.to_eager(self.model_tensors[name]()) scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]()) - # ndim alone is not enough: the FP8 group also has a 2D [out, 1] weight_scale. NVFP4 is nibble-packed uint8, E4M3 scale, one per 16. + # ndim is not enough: FP8 also has a 2D [out,1] scale. NVFP4 is packed uint8, E4M3 per 16. if scale.ndim < 2: continue if weight.dtype != torch.uint8 or scale.dtype != torch.float8_e4m3fn: From ca2f4b3887ce4294626a1cb99128ea7092ca831a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 8 Sep 2026 23:16:02 -0700 Subject: [PATCH 5/6] tighten the comments added by this change No code change: every remaining line carries a fact the code does not state. --- conversion/base.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index ffc0e8bd5516..f21423aadaae 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -538,9 +538,8 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T if (base_name + "_zero_point") in self.model_tensors: tensors_to_remove.append(base_name + "_zero_point") elif nvfp4_compressed_tensors: - # NVFP4 tensors are gone by here, so a leftover weight_scale is the FP8 group. - # Per-channel residuals only: a block group's scales are a grid needing - # block_structure expansion, which dequant_simple would misapply. + # NVFP4 is gone by here, so a leftover weight_scale is the FP8 group. Per-channel + # only: block scales are a grid dequant_simple would misapply. for group in groups.values(): if not isinstance(group, dict) or group.get("format") == "nvfp4-pack-quantized": continue From 559a2420dc7771bdc1a482d43e48088cf422b019 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 9 Sep 2026 00:17:10 -0700 Subject: [PATCH 6/6] convert: refuse residual groups dequant_simple cannot dequantize Recognising mixed-precision means the NVFP4 branch of dequant_model now runs on real inputs instead of being a bare pass, and everything the NVFP4 group does not consume lands there. The only dequantizer that branch has is dequant_simple, which is correct for one shape of input: an unpacked weight with one scale per row. The guard only checked strategy, so three things reached it that it cannot handle. A "pack-quantized" residual group is the one that matters. Its weights are nibble-packed ints needing dequant_packed, but its strategy really is "channel", so a strategy-only guard waves it through. The weight_name-not-in-model_tensors check does not catch it either, because prepare_tensors renames every .weight_packed to .weight before dequant_model runs, so by the time the branch looks, the name it is testing for absence exists. The result was a clean exit 0 and a file with packed nibbles multiplied by a scale where weights should be. unsloth/gemma-4-E2B-it-NVFP4 and unsloth/gemma-4-E4B-it-NVFP4 both ship exactly that group_2 over embed_tokens_per_layer, so this is a published model family silently mis-converting rather than a hypothetical. Before this change master refused those checkpoints outright at the multiple-config-groups raise, so widening the gate turned a hard refusal into a wrong file, which is worse than the bug it replaced. Also fixed, all in the same branch: - A weight still uint8 after _generate_nvfp4_tensors is one whose dtype or block geometry that function skipped. For most widths the shapes then fail to broadcast and you get a loud error, but a [out, 1] scale broadcasts cleanly and dequant_simple returns scaled nibbles with no error at all. Refuse it instead. - weight_config = tuple(groups.values())[0]["weights"] ran unconditionally right after the gate and is dead code for the NVFP4 branch, so a group with no "weights" key raised KeyError. Moved into the two branches that use it. - The block_structure rejection reported "a 'channel' group is not supported", naming the field that was fine and hiding the one that was not. Split into three messages that each name what actually tripped. scripts/unsloth/test_convert_nvfp4_mixed.py covers all four against the real gemma-4 group_2 config, plus controls that the supported nvfp4-plus-fp8-channel shape still converts to the right numbers and that single-group formats are untouched. It fails in seven places on the parent commit and passes here. Byte-identity re-checked with this change in: 44 conversions across six architectures and eleven quantization types, 37 byte-identical to the merge base and zero changed by this commit. unsloth/Qwen3.6-27B-NVFP4 still converts with max-abs-error 0.000e+00 against a reference dequantization from the safetensors. --- conversion/base.py | 42 ++- scripts/unsloth/test_convert_nvfp4_mixed.py | 364 ++++++++++++++++++++ 2 files changed, 400 insertions(+), 6 deletions(-) create mode 100644 scripts/unsloth/test_convert_nvfp4_mixed.py diff --git a/conversion/base.py b/conversion/base.py index f21423aadaae..000cf2d2dc0e 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -494,9 +494,9 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T if len(groups) > 1 and not nvfp4_compressed_tensors: raise NotImplementedError("Can't handle multiple config groups for compressed-tensors yet") - weight_config = tuple(groups.values())[0]["weights"] if quant_format == "float-quantized" or quant_format == "int-quantized" or quant_format == "naive-quantized": + weight_config = tuple(groups.values())[0]["weights"] block_size = weight_config.get("block_structure", None) strategy = weight_config.get("strategy") assert strategy == "channel" or strategy == "block" @@ -516,6 +516,7 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T if self._fp8_as_q8 and is_fp8: self._fp8_dequantized.add(weight_name) elif quant_format == "pack-quantized": + weight_config = tuple(groups.values())[0]["weights"] assert weight_config.get("strategy") == "group" assert weight_config.get("type", "int") == "int" num_bits = weight_config.get("num_bits") @@ -544,10 +545,27 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T if not isinstance(group, dict) or group.get("format") == "nvfp4-pack-quantized": continue residual = group.get("weights") or {} - if residual.get("strategy") != "channel" or residual.get("block_structure") is not None: + # dequant_simple is the only dequantizer reachable from here, so the + # residual group has to be one that dequant_simple is correct for: + # an unpacked weight with one scale per row. "pack-quantized" is not, + # its weights are nibble-packed ints needing dequant_packed, and + # prepare_tensors has already renamed its weight_packed to weight, so + # nothing downstream can tell. Refusing beats a silently wrong file. + group_format = group.get("format") + if group_format not in (None, "float-quantized", "int-quantized", "naive-quantized"): raise NotImplementedError( f"compressed-tensors mixed-precision with NVFP4 plus a " - f"{residual.get('strategy')!r} group is not yet supported" + f"{group_format!r} group is not yet supported" + ) + if residual.get("block_structure") is not None: + raise NotImplementedError( + f"compressed-tensors mixed-precision with NVFP4 plus a group with " + f"block_structure {residual.get('block_structure')!r} is not yet supported" + ) + if residual.get("strategy") != "channel": + raise NotImplementedError( + f"compressed-tensors mixed-precision with NVFP4 plus a " + f"{residual.get('strategy')!r} strategy group is not yet supported" ) for name in self.model_tensors.keys(): if name.endswith(".weight_scale"): @@ -557,9 +575,21 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T continue w = self.model_tensors[weight_name] s = self.model_tensors[name] - is_fp8 = False - if self._fp8_as_q8: - is_fp8 = w().dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + # _generate_nvfp4_tensors consumed every tensor it recognised as + # NVFP4, so a uint8 weight still here is one its dtype or geometry + # guard skipped. dequant_simple would multiply packed nibbles by a + # scale and write the result out as if it were real weights, which + # for some shapes broadcasts cleanly and produces no error at all. + # .dtype and .shape come off a device="meta" tensor, so this does + # not materialize anything. + w_meta = w() + if w_meta.dtype == torch.uint8: + raise NotImplementedError( + f"{weight_name!r} is still packed uint8 after NVFP4 repacking, so its " + f"block geometry is not one this converter understands " + f"(weight {tuple(w_meta.shape)}, scale {tuple(s().shape)})" + ) + is_fp8 = self._fp8_as_q8 and w_meta.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) self.model_tensors[weight_name] = lambda w=w, s=s: dequant_simple(w(), s(), None) tensors_to_remove.append(name) if is_fp8: diff --git a/scripts/unsloth/test_convert_nvfp4_mixed.py b/scripts/unsloth/test_convert_nvfp4_mixed.py new file mode 100644 index 000000000000..9d00d592d365 --- /dev/null +++ b/scripts/unsloth/test_convert_nvfp4_mixed.py @@ -0,0 +1,364 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Tests for the compressed-tensors mixed-precision NVFP4 path in conversion/base.py. +Run: python3 scripts/unsloth/test_convert_nvfp4_mixed.py + +Recognising a mixed-precision checkpoint means the NVFP4 branch of dequant_model now +runs on real inputs instead of being a bare `pass`. Everything the NVFP4 group does not +consume lands in that branch, and the only dequantizer it has is dequant_simple, which +is correct for exactly one shape of input: an unpacked weight with one scale per row. + +So the thing worth testing is not that a supported checkpoint converts. It is that an +UNSUPPORTED residual group is refused rather than quietly multiplied by a scale and +written out, because a converter's output is a file someone keeps and uses for months +and a wrong number in it never announces itself. + +Two ways the branch could be handed something dequant_simple cannot dequantize: + + 1. A "pack-quantized" residual group. Its weights are nibble-packed ints needing + dequant_packed. It passes a strategy-only guard because its strategy really is + "channel". And prepare_tensors renames every `.weight_packed` to `.weight` BEFORE + dequant_model runs, so by the time the branch sees it, the `weight_name not in + model_tensors` guard is testing a name that now exists. The guard looks correct and + is inert. That is `test_pack_quantized_residual_is_refused`, and the config it uses + is copied verbatim from unsloth/gemma-4-E2B-it-NVFP4 and -E4B-it-NVFP4, which ship + exactly this group_2 over embed_tokens_per_layer. + + 2. An NVFP4 tensor whose block geometry _generate_nvfp4_tensors skipped. It stays uint8 + and packed. For most widths the shapes then fail to broadcast and you get a loud + error, but when the scale is [out, 1] it broadcasts cleanly and dequant_simple + happily returns packed nibbles scaled as if they were weights. That is + `test_surviving_packed_uint8_is_refused`. + +dequant_model only reads self._is_nvfp4, self.model_tensors, self.hparams, +self._fp8_as_q8 and self._fp8_dequantized, so these drive it through a shim rather than +running a full conversion. No tokenizer, no network, no model download. +""" +import sys +import traceback +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +FAILS: list[str] = [] + + +def check(name: str, cond: bool, detail: str = "") -> None: + if cond: + print(f"ok {name}") + else: + print(f"FAIL {name}{': ' + detail if detail else ''}") + FAILS.append(name) + + +try: + import torch + from conversion.base import ModelBase +except ImportError as e: # pragma: no cover + print(f"SKIP: {e}. Install requirements/requirements-convert_hf_to_gguf.txt to run this.") + sys.exit(0) + + +class Shim: + """The four attributes dequant_model actually reads.""" + + def __init__(self, tensors, quant_config, is_nvfp4=True, fp8_as_q8=False): + self.model_tensors = {k: (lambda v=v: v) for k, v in tensors.items()} + self.hparams = {"quantization_config": quant_config} + self._is_nvfp4 = is_nvfp4 + self._fp8_as_q8 = fp8_as_q8 + self._fp8_dequantized = set() + + def run(self): + ModelBase.dequant_model(self) + + def get(self, name): + return self.model_tensors[name]() + + +def nvfp4_group(targets): + return { + "format": "nvfp4-pack-quantized", + "targets": targets, + "weights": { + "num_bits": 4, "type": "float", "strategy": "tensor_group", + "group_size": 16, "block_structure": None, "symmetric": True, + "scale_dtype": "torch.float8_e4m3fn", + }, + } + + +def fp8_channel_group(targets): + return { + "format": "float-quantized", + "targets": targets, + "weights": { + "num_bits": 8, "type": "float", "strategy": "channel", + "group_size": None, "block_structure": None, "symmetric": True, + }, + } + + +# Copied verbatim from unsloth/gemma-4-E2B-it-NVFP4 and unsloth/gemma-4-E4B-it-NVFP4 +# config.json, quantization_config.config_groups.group_2. Both ship this identically. +GEMMA4_PACK_QUANTIZED_GROUP = { + "format": "pack-quantized", + "input_activations": None, + "output_activations": None, + "targets": ["re:.*embed_tokens_per_layer$"], + "weights": { + "actorder": None, "block_structure": None, "dynamic": False, + "group_size": None, "num_bits": 8, "observer": "memoryless_minmax", + "observer_kwargs": {}, "scale_dtype": None, "strategy": "channel", + "symmetric": True, "type": "int", "zp_dtype": None, + }, +} + + +def residual_tensors(prefix="model.layers.0.self_attn.q_proj"): + """The state dequant_model actually sees. _generate_nvfp4_tensors has already run and + consumed every NVFP4 weight/scale pair, so what is left is the residual group only.""" + return { + f"{prefix}.weight": torch.zeros(4, 8, dtype=torch.float8_e4m3fn), + f"{prefix}.weight_scale": torch.ones(4, 1), + } + + +def raises(shim): + try: + shim.run() + except BaseException as e: # noqa: BLE001 - the type is part of what we assert + return e + return None + + +# -------------------------------------------------------------------------------------- +# 1. The defect: a pack-quantized residual group must be refused, not dequant_simple'd. +# Before the fix this returned cleanly and left embed_tokens_per_layer.weight as a +# dequant_simple lambda over nibble-packed int32. Exit code 0, wrong file. +# -------------------------------------------------------------------------------------- +def test_pack_quantized_residual_is_refused(): + tensors = residual_tensors() + tensors.update({ + # prepare_tensors has already renamed .weight_packed -> .weight, which is exactly + # why the `weight_name not in model_tensors` guard does not catch this. + "model.embed_tokens_per_layer.weight": torch.zeros(8, 16, dtype=torch.int32), + "model.embed_tokens_per_layer.weight_scale": torch.ones(8, 1, dtype=torch.bfloat16), + "model.embed_tokens_per_layer.weight_shape": torch.tensor([8, 64]), + }) + shim = Shim(tensors, { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": { + "group_0": fp8_channel_group(["re:.*self_attn\\.(q|k|v|o)_proj$"]), + "group_1": nvfp4_group(["re:.*mlp\\.(gate|up|down)_proj$"]), + "group_2": GEMMA4_PACK_QUANTIZED_GROUP, + }, + }) + err = raises(shim) + check("pack-quantized residual raises", isinstance(err, NotImplementedError), + f"got {err!r}") + check("pack-quantized message names the format", + err is not None and "pack-quantized" in str(err), f"got {err!r}") + check("pack-quantized weight was not silently dequantized", + err is not None, + "dequant_model returned cleanly, embed_tokens_per_layer is now wrong numbers") + + +# -------------------------------------------------------------------------------------- +# 2. Control: the shape this PR exists to support must still convert. A guard that +# refuses everything would pass test 1 and be useless. +# -------------------------------------------------------------------------------------- +def test_nvfp4_plus_fp8_channel_still_works(): + w = (torch.arange(32, dtype=torch.float32).reshape(4, 8) / 8.0).to(torch.float8_e4m3fn) + s = torch.tensor([[2.0], [3.0], [4.0], [5.0]]) + tensors = residual_tensors() + tensors.update({ + "model.layers.0.self_attn.q_proj.weight": w, + "model.layers.0.self_attn.q_proj.weight_scale": s, + "model.layers.0.self_attn.q_proj.input_scale": torch.tensor(1.0), + "model.layers.0.self_attn.k_scale": torch.tensor(1.0), + "model.layers.0.self_attn.v_scale": torch.tensor(1.0), + }) + shim = Shim(tensors, { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": { + "group_0": fp8_channel_group(["re:.*self_attn\\.(q|k|v|o)_proj$"]), + "group_1": nvfp4_group(["re:.*mlp\\.(gate|up|down)_proj$"]), + }, + }, fp8_as_q8=True) + err = raises(shim) + check("nvfp4 + fp8-channel does not raise", err is None, f"got {err!r}") + if err is not None: + return + got = shim.get("model.layers.0.self_attn.q_proj.weight") + check("fp8 residual dequantized to the right numbers", + torch.allclose(got, w.float() * s), f"got {got}") + for sidecar in ("weight_scale", "input_scale"): + check(f"{sidecar} sidecar dropped", + f"model.layers.0.self_attn.q_proj.{sidecar}" not in shim.model_tensors) + for sidecar in ("k_scale", "v_scale"): + check(f"{sidecar} sidecar dropped", + f"model.layers.0.self_attn.{sidecar}" not in shim.model_tensors) + check("fp8 weight recorded for --fp8-as-q8", + "model.layers.0.self_attn.q_proj.weight" in shim._fp8_dequantized) + + +# -------------------------------------------------------------------------------------- +# 3. A residual group with a block_structure is refused, and says so. dequant_simple +# would apply a 2-D grid of scales as if it were per-row. +# -------------------------------------------------------------------------------------- +def test_block_structure_residual_is_refused(): + group = fp8_channel_group(["re:.*self_attn\\.(q|k|v|o)_proj$"]) + group["weights"]["block_structure"] = [128, 128] + shim = Shim(residual_tensors(), { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": {"group_0": group, "group_1": nvfp4_group(["re:.*mlp.*$"])}, + }) + err = raises(shim) + check("block_structure residual raises", isinstance(err, NotImplementedError), + f"got {err!r}") + # The message used to say "a 'channel' group is not supported", naming the field + # that was fine and hiding the field that was not. + check("block_structure message names block_structure", + err is not None and "block_structure" in str(err), f"got {err!r}") + + +def test_non_channel_strategy_residual_is_refused(): + for strategy in ("tensor", "group", "token", None): + group = fp8_channel_group(["re:.*self_attn.*$"]) + group["weights"]["strategy"] = strategy + shim = Shim(residual_tensors(), { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": {"group_0": group, "group_1": nvfp4_group(["re:.*mlp.*$"])}, + }) + err = raises(shim) + check(f"strategy={strategy!r} residual raises", + isinstance(err, NotImplementedError), f"got {err!r}") + + +# -------------------------------------------------------------------------------------- +# 4. Malformed and older exports must fail as NotImplementedError, never as a KeyError +# or AttributeError from an unguarded dict access. +# -------------------------------------------------------------------------------------- +def test_missing_weights_key_is_not_a_keyerror(): + for weights in (None, {}): + group = {"format": "float-quantized", "targets": ["re:.*self_attn.*$"], "weights": weights} + shim = Shim(residual_tensors(), { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": {"group_0": group, "group_1": nvfp4_group(["re:.*mlp.*$"])}, + }) + err = raises(shim) + check(f"weights={weights!r} raises NotImplementedError", + isinstance(err, NotImplementedError), f"got {type(err).__name__}: {err}") + + # weights absent entirely. `weight_config = tuple(groups.values())[0]["weights"]` used + # to run unconditionally just after the gate and KeyError here, even though the NVFP4 + # branch never uses weight_config. + group = {"format": "float-quantized", "targets": ["re:.*self_attn.*$"]} + shim = Shim(residual_tensors(), { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": {"group_0": group, "group_1": nvfp4_group(["re:.*mlp.*$"])}, + }) + err = raises(shim) + check("absent weights key raises NotImplementedError not KeyError", + isinstance(err, NotImplementedError), f"got {type(err).__name__}: {err}") + + +# -------------------------------------------------------------------------------------- +# 5. An NVFP4 tensor _generate_nvfp4_tensors skipped stays packed uint8. With a [out, 1] +# scale it broadcasts cleanly, so this is the one shape where the old code produced a +# file with no error at all. +# -------------------------------------------------------------------------------------- +def test_surviving_packed_uint8_is_refused(): + shim = Shim({ + "model.layers.0.mlp.down_proj.weight": torch.full((4, 8), 0x42, dtype=torch.uint8), + "model.layers.0.mlp.down_proj.weight_scale": torch.ones(4, 1).to(torch.float8_e4m3fn), + }, { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": { + "group_0": fp8_channel_group(["re:.*self_attn.*$"]), + "group_1": nvfp4_group(["re:.*mlp.*$"]), + }, + }) + err = raises(shim) + check("surviving packed uint8 raises", isinstance(err, NotImplementedError), + f"got {err!r}") + check("packed uint8 message names the dtype", + err is not None and "uint8" in str(err), f"got {err!r}") + + +# -------------------------------------------------------------------------------------- +# 6. The gate itself. any() must not pull a mixed-precision checkpoint with no NVFP4 +# group into the NVFP4 path. +# -------------------------------------------------------------------------------------- +def test_mixed_precision_without_nvfp4_still_refused(): + shim = Shim({ + "model.layers.0.self_attn.q_proj.weight": torch.zeros(4, 8, dtype=torch.float8_e4m3fn), + "model.layers.0.self_attn.q_proj.weight_scale": torch.ones(4, 1), + }, { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "config_groups": { + "group_0": fp8_channel_group(["re:.*self_attn.*$"]), + "group_1": fp8_channel_group(["re:.*mlp.*$"]), + }, + }, is_nvfp4=False) + err = raises(shim) + check("mixed-precision with no NVFP4 group is refused", + isinstance(err, NotImplementedError) and "multiple config groups" in str(err), + f"got {err!r}") + + +# -------------------------------------------------------------------------------------- +# 7. Single-group formats must be untouched by all of the above. +# -------------------------------------------------------------------------------------- +def test_single_group_float_quantized_unaffected(): + w = torch.zeros(4, 8, dtype=torch.float8_e4m3fn) + s = torch.full((4, 1), 2.0) + shim = Shim({ + "model.layers.0.self_attn.q_proj.weight": w, + "model.layers.0.self_attn.q_proj.weight_scale": s, + }, { + "quant_method": "compressed-tensors", + "format": "float-quantized", + "config_groups": {"group_0": fp8_channel_group(["re:.*self_attn.*$"])}, + }, is_nvfp4=False, fp8_as_q8=True) + err = raises(shim) + check("single-group float-quantized does not raise", err is None, f"got {err!r}") + if err is None: + check("single-group float-quantized dequantizes correctly", + torch.allclose(shim.get("model.layers.0.self_attn.q_proj.weight"), w.float() * s)) + + +if __name__ == "__main__": + for fn in ( + test_pack_quantized_residual_is_refused, + test_nvfp4_plus_fp8_channel_still_works, + test_block_structure_residual_is_refused, + test_non_channel_strategy_residual_is_refused, + test_missing_weights_key_is_not_a_keyerror, + test_surviving_packed_uint8_is_refused, + test_mixed_precision_without_nvfp4_still_refused, + test_single_group_float_quantized_unaffected, + ): + print(f"-- {fn.__name__}") + try: + fn() + except Exception: + traceback.print_exc() + FAILS.append(fn.__name__) + + print() + if FAILS: + print(f"{len(FAILS)} failure(s): {', '.join(FAILS)}") + sys.exit(1) + print("all passed")