From 9702f30fd6705a7b900745fe3918c378fe0e4088 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 11:47:24 -0700 Subject: [PATCH 01/14] nvfp4: wip --- gptqmodel/looper/weight_only_looper.py | 18 +- gptqmodel/looper/weight_only_processor.py | 15 +- gptqmodel/models/auto.py | 2 + gptqmodel/models/base.py | 3 + gptqmodel/nn_modules/qlinear/fp4.py | 339 +++++++++++++++- gptqmodel/quantization/__init__.py | 1 + gptqmodel/quantization/config.py | 161 +++++++- gptqmodel/quantization/gptq.py | 6 +- gptqmodel/quantization/quantizer.py | 50 ++- gptqmodel/utils/backend.py | 6 + gptqmodel/utils/model.py | 3 +- gptqmodel/utils/structure.py | 9 + .../minimax_m3_nvfp4/debug_expert_capture.py | 93 +++++ scripts/minimax_m3_nvfp4/p03_load_dryrun.py | 116 ++++++ scripts/minimax_m3_nvfp4/p1_pilot_layer3.py | 123 ++++++ .../minimax_m3_nvfp4/p1b_expert_coverage.py | 167 ++++++++ tests/test_nvfp4_export.py | 369 ++++++++++++++++++ 17 files changed, 1458 insertions(+), 23 deletions(-) create mode 100644 scripts/minimax_m3_nvfp4/debug_expert_capture.py create mode 100644 scripts/minimax_m3_nvfp4/p03_load_dryrun.py create mode 100644 scripts/minimax_m3_nvfp4/p1_pilot_layer3.py create mode 100644 scripts/minimax_m3_nvfp4/p1b_expert_coverage.py create mode 100644 tests/test_nvfp4_export.py diff --git a/gptqmodel/looper/weight_only_looper.py b/gptqmodel/looper/weight_only_looper.py index 87b9938ac..77e698e2e 100644 --- a/gptqmodel/looper/weight_only_looper.py +++ b/gptqmodel/looper/weight_only_looper.py @@ -30,7 +30,7 @@ from ..models import BaseQModel from ..models._const import CPU, SUPPORTS_MODULE_TYPES from ..nn_modules.converter import MODULE_CONVERTER_MAP -from ..quantization.config import BitsAndBytesConfig, FP8Config, GGUFConfig, RTNConfig, VramStrategy +from ..quantization.config import BitsAndBytesConfig, FP8Config, GGUFConfig, NVFP4Config, RTNConfig, VramStrategy from ..utils import has_gil_disabled from ..utils.device import get_device from ..utils.device_telemetry import emit_device_telemetry @@ -414,7 +414,7 @@ def _quantize_named_module( self, named: NamedModule, target_device: torch.device, - ) -> Tuple[NamedModule, Optional[RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]]: + ) -> Tuple[NamedModule, Optional[RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]]: """Run one module's weight-only quantization on its assigned device.""" self._prepare_named_module_for_quantization(named, target_device) @@ -427,7 +427,7 @@ def _quantize_named_module( def _finalize_quantized_module( self, named: NamedModule, - active_qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig, + active_qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig, ) -> str: """Move a quantized module back to CPU, pack it, and optionally offload it.""" @@ -477,7 +477,7 @@ def _finalize_quantized_module( def _finalize_target_device( self, - active_qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig, + active_qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig, ) -> torch.device: """Resolve the worker device for one finalize task.""" @@ -488,7 +488,7 @@ def _finalize_target_device( def _finalize_subset_modules( self, - quantized_modules: List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]], + quantized_modules: List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]], ) -> None: """Finalize one subset, using the device pool when multiple finalize targets exist.""" @@ -576,10 +576,10 @@ def _advance_finalize_progress(named: NamedModule, module_label: str) -> None: def _quantize_subset_modules( self, named_modules: List[NamedModule], - ) -> List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]]: + ) -> List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]]: """Quantize one subset, using multiple devices when available.""" - results_by_name: Dict[str, Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]] = {} + results_by_name: Dict[str, Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]] = {} task_specs = [ (named, self._assign_quant_device_for_module(named, CPU)) for named in named_modules @@ -669,10 +669,10 @@ def _offload_quantized_module(self, module: NamedModule) -> None: def loop(self, **kwargs): """Quantize layers directly from weights without calibration forwards.""" quant_config = self.gptq_model.quantize_config - if not isinstance(quant_config, (RTNConfig, GGUFConfig, FP8Config, BitsAndBytesConfig)): + if not isinstance(quant_config, (RTNConfig, GGUFConfig, FP8Config, NVFP4Config, BitsAndBytesConfig)): raise NotImplementedError( "Weight-only looper only supports `RTNConfig`, `GGUFConfig`, " - "`FP8Config`, and `BitsAndBytesConfig` today." + "`FP8Config`, `NVFP4Config`, and `BitsAndBytesConfig` today." ) if quant_config.lm_head: diff --git a/gptqmodel/looper/weight_only_processor.py b/gptqmodel/looper/weight_only_processor.py index 31b08abe7..1e4615aef 100644 --- a/gptqmodel/looper/weight_only_processor.py +++ b/gptqmodel/looper/weight_only_processor.py @@ -29,6 +29,7 @@ FP8Config, GGUFConfig, METHOD, + NVFP4Config, RTNConfig, clone_weight_only_config_for_module, resolve_quant_format, @@ -48,9 +49,9 @@ class WeightOnlyProcessor(LoopProcessor): def __init__( self, tokenizer, - qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig, + qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig, ): - """Initializes a weight-only processor for RTN, GGUF, FP8, or BitsAndBytes.""" + """Initializes a weight-only processor for RTN, GGUF, FP8, NVFP4, or BitsAndBytes.""" super().__init__( tokenizer=tokenizer, @@ -74,10 +75,10 @@ def is_skipped(self, module: NamedModule) -> bool: return self.qcfg.dynamic_get(layer_name=module.full_name) is False @staticmethod - def _uses_direct_pack(qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig) -> bool: + def _uses_direct_pack(qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig) -> bool: """Returns whether the method packs directly from the original dense weights.""" - return qcfg.method in {METHOD.GGUF, METHOD.FP8, METHOD.BITSANDBYTES} + return qcfg.method in {METHOD.GGUF, METHOD.FP8, METHOD.NVFP4, METHOD.BITSANDBYTES} def _update_logged_loss(self, module: NamedModule, avg_loss: str) -> None: """Backfills the logged loss field after late dequant-error measurement.""" @@ -93,7 +94,7 @@ def quantize_module( module: NamedModule, *, device: Optional[torch.device] = None, - ) -> Optional[RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]: + ) -> Optional[RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]: """Clones per-module config, quantizes weights, and logs the result.""" qcfg_clone = clone_weight_only_config_for_module(self.qcfg, module.full_name) @@ -149,7 +150,7 @@ def submodule_finalize( module: NamedModule, model: BaseQModel, *, - qcfg: Optional[RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig] = None, + qcfg: Optional[RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig] = None, **kwargs, ): """Creates and packs the final quantized module into the model graph.""" @@ -268,6 +269,8 @@ def name(self) -> str: return "weight_only_gguf" if self.qcfg.method == METHOD.FP8: return "weight_only_fp8" + if self.qcfg.method == METHOD.NVFP4: + return "weight_only_nvfp4" if self.qcfg.method == METHOD.BITSANDBYTES: return "weight_only_bitsandbytes" return "weight_only_rtn" diff --git a/gptqmodel/models/auto.py b/gptqmodel/models/auto.py index 478afc2f9..ed2699749 100644 --- a/gptqmodel/models/auto.py +++ b/gptqmodel/models/auto.py @@ -730,6 +730,8 @@ def export(model_id_or_path: str, target_path: str, format: str, trust_remote_co backend = BACKEND.PAROQUANT_CUDA elif normalized_method == METHOD.FP8.value: backend = BACKEND.FP8_TORCH + elif normalized_method == METHOD.NVFP4.value: + backend = BACKEND.NVFP4_TORCH elif normalized_method == METHOD.EXL3.value: backend = BACKEND.EXL3_TORCH else: diff --git a/gptqmodel/models/base.py b/gptqmodel/models/base.py index e045547a3..aab73a22e 100644 --- a/gptqmodel/models/base.py +++ b/gptqmodel/models/base.py @@ -896,6 +896,9 @@ def quantize( preferred_backend = BACKEND.AUTO elif self.quantize_config.method == METHOD.FP8: preferred_backend = BACKEND.FP8_TORCH + elif export_quant_method == METHOD.NVFP4 or format_code == FORMAT.NVFP4: + # Covers both weight-only NVFP4 and GPTQ-on-NVFP4-grid exports. + preferred_backend = BACKEND.NVFP4_TORCH elif self.quantize_config.method == METHOD.BITSANDBYTES: preferred_backend = BACKEND.BITSANDBYTES else: diff --git a/gptqmodel/nn_modules/qlinear/fp4.py b/gptqmodel/nn_modules/qlinear/fp4.py index 0eb6a0df8..25875e88f 100644 --- a/gptqmodel/nn_modules/qlinear/fp4.py +++ b/gptqmodel/nn_modules/qlinear/fp4.py @@ -9,11 +9,25 @@ import torch.nn as nn import torch.nn.functional as F +from ...adapter.adapter import Adapter, Lora +from ...models._const import DEVICE, PLATFORM +from ...quantization import FORMAT, METHOD +from ...quantization.config import NVFP4_WEIGHT_BLOCK_SIZE +from ...quantization.dtype import dequantize_f4_e2m1, device_supports_native_fp4 +from ...utils.backend import BACKEND +from . import WeightOnlyQuantLinear +from .fp8 import _weight_to_matrix +from .gguf import _apply_optional_smoother + try: - from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor + from torchao.prototype.mx_formats.kernels import f32_to_f4_unpacked, pack_uint4 + from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor, nvfp4_quantize except Exception: NVFP4Tensor = None + nvfp4_quantize = None + f32_to_f4_unpacked = None + pack_uint4 = None class TorchFP4Linear(nn.Module): @@ -61,3 +75,326 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if isinstance(bias, torch.Tensor) and bias.dtype != x.dtype: bias = bias.to(dtype=x.dtype) return F.linear(x, self._native_weight(), bias) + + +def quantize_nvfp4_weight( + weight: torch.Tensor, + *, + weight_block_size: int = NVFP4_WEIGHT_BLOCK_SIZE, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize one dense 2D weight to NVFP4 packed bytes plus per-block E4M3 scales.""" + + if nvfp4_quantize is None: + raise RuntimeError("NVFP4 quantization requires torchao with NVFP4 support.") + if weight.ndim != 2: + raise ValueError(f"NVFP4 quantization expects a 2D weight matrix, got shape {tuple(weight.shape)}.") + if int(weight_block_size) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError(f"NVFP4 quantization requires `weight_block_size={NVFP4_WEIGHT_BLOCK_SIZE}`.") + if weight.shape[1] % NVFP4_WEIGHT_BLOCK_SIZE != 0: + raise ValueError( + f"NVFP4 quantization expects in_features divisible by {NVFP4_WEIGHT_BLOCK_SIZE}, " + f"got shape {tuple(weight.shape)}." + ) + + weight = weight.to(device="cpu", dtype=torch.float32).contiguous() + scales, packed = nvfp4_quantize(weight, block_size=NVFP4_WEIGHT_BLOCK_SIZE) + # Store packed nibbles as uint8 so safetensors serialization matches the + # ModelOpt/torchao NVFP4 checkpoint layout regardless of torch fp4 dtype support. + if packed.dtype is not torch.uint8: + packed = packed.view(torch.uint8) + return packed.contiguous(), scales.contiguous() + + +class TorchNVFP4Linear(WeightOnlyQuantLinear): + SUPPORTS_BACKENDS = [BACKEND.NVFP4_TORCH] + # METHOD.GPTQ is supported with FORMAT.NVFP4: the GPTQ loop quantizes on the + # NVFP4 grid and this kernel stores/executes the packed result. + SUPPORTS_METHODS = [METHOD.NVFP4, METHOD.GPTQ] + SUPPORTS_FORMATS = {FORMAT.NVFP4: 15} + SUPPORTS_BITS = [4] + SUPPORTS_GROUP_SIZE = [-1, NVFP4_WEIGHT_BLOCK_SIZE] + SUPPORTS_SYM = [True] + SUPPORTS_DESC_ACT = [False] + SUPPORTS_SHARDS = True + SUPPORTS_TRAINING = False + SUPPORTS_AUTO_PADDING = False + SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [NVFP4_WEIGHT_BLOCK_SIZE] + SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [1] + SUPPORTS_DEVICES = [DEVICE.CPU, DEVICE.CUDA, DEVICE.ROCM, DEVICE.XPU, DEVICE.MPS] + SUPPORTS_PLATFORM = [PLATFORM.ALL] + SUPPORTS_PACK_DTYPES = [torch.int8, torch.int16, torch.int32, torch.int64] + SUPPORTS_ADAPTERS = [Lora] + SUPPORTS_DTYPES = [torch.float16, torch.bfloat16, torch.float32] + + QUANT_TYPE = "nvfp4" + + def __init__( + self, + bits: int, + group_size: int, + sym: bool, + desc_act: bool, + in_features: int, + out_features: int, + bias: bool = False, + pack_dtype: torch.dtype = torch.int32, + adapter: Adapter = None, + register_buffers: bool = True, + weight_block_size: int = NVFP4_WEIGHT_BLOCK_SIZE, + **kwargs, + ): + if int(weight_block_size) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"TorchNVFP4Linear requires `weight_block_size={NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + self.weight_block_size = NVFP4_WEIGHT_BLOCK_SIZE + self._native_linear_hard_disabled = False + + super().__init__( + bits=bits, + in_features=in_features, + out_features=out_features, + bias=bias, + backend=kwargs.pop("backend", BACKEND.NVFP4_TORCH), + adapter=adapter, + register_buffers=False, + pack_dtype=pack_dtype, + **kwargs, + ) + + if register_buffers: + self._allocate_buffers(bias=bias) + + @classmethod + def validate_once(cls): + if NVFP4Tensor is None or nvfp4_quantize is None: + return False, RuntimeError("TorchNVFP4Linear requires torchao with NVFP4 support.") + return True, None + + def smooth_block_size(self) -> int: + return self.weight_block_size + + def _allocate_buffers(self, *, bias: bool) -> None: + weight = torch.zeros((self.out_features, self.in_features // 2), dtype=torch.uint8) + scale = torch.zeros( + (self.out_features, self.in_features // self.weight_block_size), + dtype=torch.float8_e4m3fn, + ) + + if "weight" in self._buffers: + self.weight = weight + else: + self.register_buffer("weight", weight) + + if "weight_scale" in self._buffers: + self.weight_scale = scale + else: + self.register_buffer("weight_scale", scale) + + if bias: + bias_tensor = torch.zeros(self.out_features, dtype=torch.float16) + if "bias" in self._buffers: + self.bias = bias_tensor + else: + self.register_buffer("bias", bias_tensor) + else: + self.bias = None + + def list_buffers(self): + buffers = [] + if hasattr(self, "weight") and self.weight is not None: + buffers.append(self.weight) + if hasattr(self, "weight_scale") and self.weight_scale is not None: + buffers.append(self.weight_scale) + if hasattr(self, "bias") and self.bias is not None: + buffers.append(self.bias) + return buffers + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"bias={self.bias is not None}, weight_block_size={self.weight_block_size}" + ) + + def _weight_to_matrix(self, linear: nn.Module) -> torch.Tensor: + return _weight_to_matrix(linear) + + def pack(self, linear: nn.Module, scales: torch.Tensor, zeros: torch.Tensor, g_idx: torch.Tensor = None): + self.pack_original(linear=linear, scales=scales, zeros=zeros, g_idx=g_idx) + + def pack_block( + self, + linear: nn.Module, + scales: torch.Tensor, + zeros: torch.Tensor, + g_idx: torch.Tensor = None, + block_in: int = 8192, + workers: int = 1, + ): + del block_in, workers + self.pack_original(linear=linear, scales=scales, zeros=zeros, g_idx=g_idx) + + def pack_gpu( + self, + linear: nn.Module, + scales: torch.Tensor, + zeros: torch.Tensor, + g_idx: torch.Tensor = None, + *, + block_in: int = 8192, + device: torch.device | None = None, + ): + del block_in, device + self.pack_original(linear=linear, scales=scales, zeros=zeros, g_idx=g_idx) + + def _pack_with_gptq_scales(self, weight: torch.Tensor, scales: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a weight already on the NVFP4 grid using the loop-produced group scales. + + GPTQ error feedback mutates columns after each group scale is fixed, so the + final weight must be encoded against those exact scales — re-deriving scales + from the weight (RTN-style) would land some blocks on a different grid. + """ + + if f32_to_f4_unpacked is None or pack_uint4 is None: + raise RuntimeError("NVFP4 packing requires torchao with NVFP4 support.") + + scale = scales.to(device="cpu", dtype=torch.float32) + expected_groups = self.in_features // self.weight_block_size + if scale.shape != (self.out_features, expected_groups): + raise ValueError( + f"NVFP4 pack expects scales of shape {(self.out_features, expected_groups)}, " + f"got {tuple(scale.shape)}." + ) + + expanded = scale.repeat_interleave(self.weight_block_size, dim=1) + scaled = torch.clamp(weight / expanded, -6.0, 6.0) + codes = f32_to_f4_unpacked(scaled.contiguous()) + packed = pack_uint4(codes) + if packed.dtype is not torch.uint8: + packed = packed.view(torch.uint8) + return packed.contiguous(), scale.to(torch.float8_e4m3fn).contiguous() + + @torch.inference_mode() + def pack_original( + self, + linear: nn.Module, + scales: torch.Tensor, + zeros: torch.Tensor, + g_idx: torch.Tensor = None, + *, + smooth=None, + ): + del zeros, g_idx + + weight = self._weight_to_matrix(linear).to(device="cpu", dtype=torch.float32) + if isinstance(scales, torch.Tensor) and scales.numel() > 0: + qweight, weight_scale = self._pack_with_gptq_scales(weight, scales) + else: + weight = _apply_optional_smoother( + weight, + smooth=smooth, + group_size=self.smooth_block_size(), + ) + qweight, weight_scale = quantize_nvfp4_weight( + weight, + weight_block_size=self.weight_block_size, + ) + + if "weight" in self._buffers: + self.weight = qweight + else: + self.register_buffer("weight", qweight) + + if "weight_scale" in self._buffers: + self.weight_scale = weight_scale + else: + self.register_buffer("weight_scale", weight_scale) + + if linear.bias is not None: + bias = linear.bias.detach().to(device="cpu", dtype=torch.float16) + if "bias" in self._buffers: + self.bias = bias + else: + self.register_buffer("bias", bias) + else: + self.bias = None + + self._native_linear_hard_disabled = False + + def dequantize_weight( + self, + *, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + target_device = self.weight.device if device is None else torch.device(device) + target_dtype = torch.float32 if dtype is None else dtype + + # dequantize_f4_e2m1 handles both the fused CPU path (bf16/fp16 targets) + # and the reference torchao path for other targets/devices. + dequant_dtype = target_dtype if target_dtype in (torch.bfloat16, torch.float16) else torch.bfloat16 + source = self.weight if self.weight.device == target_device else self.weight.to(device=target_device) + scale = ( + self.weight_scale + if self.weight_scale.device == target_device + else self.weight_scale.to(device=target_device) + ) + weight = dequantize_f4_e2m1( + source, + scale=scale, + axis=None, + target_dtype=dequant_dtype, + ) + if weight.dtype != target_dtype: + weight = weight.to(dtype=target_dtype) + return weight.transpose(0, 1).contiguous() + + def _native_weight(self) -> "NVFP4Tensor": + if NVFP4Tensor is None: + raise RuntimeError("TorchNVFP4Linear requires torchao NVFP4Tensor support.") + return NVFP4Tensor( + self.weight, + self.weight_scale, + block_size=self.weight_block_size, + orig_dtype=torch.bfloat16, + ) + + def _can_use_native_linear(self, x_flat: torch.Tensor) -> bool: + return ( + not self._native_linear_hard_disabled + and NVFP4Tensor is not None + and x_flat.device.type == "cuda" + and x_flat.dtype is torch.bfloat16 + and device_supports_native_fp4(x_flat.device) + ) + + def _forward_dequant_matmul(self, x_flat: torch.Tensor) -> torch.Tensor: + weight = self.dequantize_weight(device=x_flat.device, dtype=x_flat.dtype) + return torch.matmul(x_flat, weight) + + def forward(self, x: torch.Tensor): + original_shape = x.shape[:-1] + (self.out_features,) + x_flat = x.reshape(-1, x.shape[-1]) + + if self._can_use_native_linear(x_flat): + try: + output = F.linear(x_flat, self._native_weight(), None) + except Exception: + self._native_linear_hard_disabled = True + output = self._forward_dequant_matmul(x_flat) + else: + output = self._forward_dequant_matmul(x_flat) + + if self.bias is not None: + bias = self.bias + if bias.device != output.device or bias.dtype != output.dtype: + bias = bias.to(device=output.device, dtype=output.dtype) + output = output + bias + + if self.adapter: + output = self.adapter.apply(x=x_flat, out=output) + + return output.reshape(original_shape) + + +__all__ = ["TorchFP4Linear", "TorchNVFP4Linear", "quantize_nvfp4_weight"] diff --git a/gptqmodel/quantization/__init__.py b/gptqmodel/quantization/__init__.py index c2261e779..d1fa8c5f6 100644 --- a/gptqmodel/quantization/__init__.py +++ b/gptqmodel/quantization/__init__.py @@ -27,6 +27,7 @@ GPTAQConfig, GPTQConfig, HessianConfig, + NVFP4Config, ParoConfig, PreProcessorCode, PreProcessorConfig, diff --git a/gptqmodel/quantization/config.py b/gptqmodel/quantization/config.py index 658c59126..32c18b238 100644 --- a/gptqmodel/quantization/config.py +++ b/gptqmodel/quantization/config.py @@ -114,6 +114,7 @@ class FORMAT(str, Enum): GPTQ_P = "gptq_p" GGUF = "gguf" FP8 = "fp8" + NVFP4 = "nvfp4" BITSANDBYTES = "bitsandbytes" MARLIN = "marlin" BITBLAS = "bitblas" @@ -134,6 +135,7 @@ class METHOD(str, Enum): GPTQ = "gptq" GGUF = "gguf" FP8 = "fp8" + NVFP4 = "nvfp4" BITSANDBYTES = "bitsandbytes" QQQ = "qqq" AWQ = "awq" @@ -652,6 +654,8 @@ def resolve_quant_format( return FORMAT.GGUF if method == METHOD.FP8: return FORMAT.FP8 + if method == METHOD.NVFP4: + return FORMAT.NVFP4 if method == METHOD.BITSANDBYTES: return FORMAT.BITSANDBYTES if method == METHOD.EXL3: @@ -1404,10 +1408,14 @@ def to_dict(self) -> Dict[str, Any]: FORMAT.GPTQ_P, FORMAT.MARLIN, FORMAT.BITBLAS, + FORMAT.NVFP4, }, METHOD.FP8: { FORMAT.FP8, }, + METHOD.NVFP4: { + FORMAT.NVFP4, + }, METHOD.BITSANDBYTES: { FORMAT.BITSANDBYTES, }, @@ -1439,6 +1447,7 @@ def to_dict(self) -> Dict[str, Any]: FORMAT.GPTQ_P, FORMAT.MARLIN, FORMAT.BITBLAS, + FORMAT.NVFP4, ) AWQ_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( FORMAT.GEMM, @@ -1461,6 +1470,22 @@ def to_dict(self) -> Dict[str, Any]: FP8_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( FORMAT.FP8, ) +NVFP4_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( + FORMAT.NVFP4, +) +# NVFP4 fixes 16 FP4 values per FP8 E4M3 scale along the in-feature axis. +NVFP4_WEIGHT_BLOCK_SIZE = 16 + + +def _normalize_nvfp4_format(value: Optional[Union[str, FORMAT]]) -> str: + """Normalize accepted NVFP4 format aliases (`nvfp4`, `fp4`) to `nvfp4`.""" + + if value is None: + return FORMAT.NVFP4.value + normalized = str(getattr(value, "value", value)).strip().lower() + if normalized in {"nvfp4", "fp4"}: + return FORMAT.NVFP4.value + raise ValueError(f"NVFP4Config: unsupported `format`: `{value}`. Expected `nvfp4`.") BITSANDBYTES_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( FORMAT.BITSANDBYTES, ) @@ -1485,6 +1510,7 @@ def to_dict(self) -> Dict[str, Any]: FORMAT.GPTQ_V2: METHOD.GPTQ, FORMAT.GPTQ_P: METHOD.GPTQ, FORMAT.FP8: METHOD.FP8, + FORMAT.NVFP4: METHOD.NVFP4, FORMAT.BITSANDBYTES: METHOD.BITSANDBYTES, FORMAT.EXL3: METHOD.EXL3, FORMAT.GGUF: METHOD.GGUF, @@ -2258,6 +2284,17 @@ def _normalize_fp8_kwargs(payload: Dict[str, Any]) -> Dict[str, Any]: return normalized +def _normalize_nvfp4_kwargs(payload: Dict[str, Any]) -> Dict[str, Any]: + normalized = dict(payload) + weight_only = normalized.pop("weight_only", None) + + if "smoother" not in normalized and "smooth" not in normalized: + normalized["smoother"] = _extract_weight_only_smooth(weight_only) + + normalized[FORMAT_FIELD_CODE] = _normalize_nvfp4_format(normalized.get(FORMAT_FIELD_CODE)) + return normalized + + def _normalize_bitsandbytes_kwargs(payload: Dict[str, Any]) -> Dict[str, Any]: normalized = dict(payload) weight_only = normalized.pop("weight_only", None) @@ -2303,6 +2340,8 @@ def _normalize_quantize_config_payload_for_target_cls(target_cls, payload: Dict[ expected_method = METHOD.AWQ elif target_cls is FP8Config: expected_method = METHOD.FP8 + elif target_cls is NVFP4Config: + expected_method = METHOD.NVFP4 elif target_cls is BitsAndBytesConfig: expected_method = METHOD.BITSANDBYTES elif target_cls is EXL3Config: @@ -2384,6 +2423,8 @@ def _prepare_target_quantize_config_kwargs(target_cls, payload: Dict[str, Any]) normalized = _normalize_gguf_kwargs(normalized) elif target_cls is FP8Config: normalized = _normalize_fp8_kwargs(normalized) + elif target_cls is NVFP4Config: + normalized = _normalize_nvfp4_kwargs(normalized) elif target_cls is BitsAndBytesConfig: normalized = _normalize_bitsandbytes_kwargs(normalized) return _filter_quantize_config_payload_for_target_cls(target_cls, normalized) @@ -3249,6 +3290,26 @@ def __post_init__(self): if self.act_group_aware and self.desc_act: raise ValueError("QuantizeConfig:: `act_group_aware` == `True` requires `desc_act` == `False`.") + if self.format == FORMAT.NVFP4: + # NVFP4 export stores contiguous 16-wide FP4 blocks with E4M3 scales, so the + # GPTQ loop must quantize on that exact grid without column reordering. + if quant_bits_width(self.bits) != 4: + raise ValueError("QuantizeConfig: `format=nvfp4` requires `bits=4`.") + if self.group_size != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"QuantizeConfig: `format=nvfp4` requires `group_size={NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + if not self.sym: + raise ValueError("QuantizeConfig: `format=nvfp4` requires `sym=True`.") + if self.desc_act: + raise ValueError("QuantizeConfig: `format=nvfp4` does not support `desc_act=True`.") + if self.act_group_aware: + if act_group_aware_user_value: + raise ValueError("QuantizeConfig: `format=nvfp4` does not support `act_group_aware=True`.") + self.act_group_aware = False + if float(self.mse or 0.0) > 0.0: + raise ValueError("QuantizeConfig: `format=nvfp4` does not support `mse` scale search.") + def _resolve_activation_ordering( self, desc_act_user_value: Optional[bool], @@ -3651,6 +3712,90 @@ def _update_output_payload(self, out: Dict[str, Any]) -> None: def uses_weight_only_lifecycle(self) -> bool: return True +@dataclass +class NVFP4Config(PreProcessorConfig): + """Weight-only NVFP4 (FP4 E2M1 + per-16-block FP8 E4M3 scales) export config.""" + + bits: int = field(default=4, metadata={"choices": [4]}) + method: METHOD = field(default=METHOD.NVFP4) + format: Optional[str] = field(default=FORMAT.NVFP4) + group_size: int = field(default=-1) + desc_act: Optional[bool] = field(default=False) + sym: bool = field(default=True) + # NVFP4 fixes the scaling block to 16 values along in-features. + weight_block_size: int = field(default=NVFP4_WEIGHT_BLOCK_SIZE) + + def _resolve_checkpoint_format(self) -> FORMAT: + self.format = FORMAT.NVFP4 + return FORMAT.NVFP4 + + def allowed_quant_methods(self) -> Tuple[METHOD, ...]: + return (METHOD.NVFP4,) + + def supported_export_formats(self) -> Tuple[FORMAT, ...]: + return NVFP4_EXPORT_FORMATS + + def default_desc_act(self) -> bool: + return False + + def __post_init__(self): + self._normalize_preprocessor_state() + super().__post_init__() + + if self.bits != 4: + raise ValueError("NVFP4Config: `bits` must be `4`.") + + if self.method != METHOD.NVFP4: + raise ValueError("NVFP4Config: `method` must be `nvfp4`.") + + self.group_size = -1 + self.desc_act = False + self.sym = True + self.format = _normalize_nvfp4_format(self.format) + + if int(self.weight_block_size) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"NVFP4Config: `weight_block_size` must be `{NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + self.weight_block_size = NVFP4_WEIGHT_BLOCK_SIZE + + if self.dynamic is not None: + self.dynamic = { + **{k: v for k, v in self.dynamic.items() if k.startswith('-')}, + **{k: v for k, v in self.dynamic.items() if not k.startswith('-')}, + } + for layer, layer_dict in self.dynamic.items(): + self._normalize_dynamic_layer_config(layer, layer_dict) + + def _normalize_dynamic_layer_config( + self, + layer_name: str, + layer_dict: Dict[str, Any], + ) -> None: + if "bits" in layer_dict and int(layer_dict["bits"]) != 4: + raise ValueError(f"NVFP4Config: layer `{layer_name}` only supports 4-bit NVFP4 weights.") + if "group_size" in layer_dict and layer_dict["group_size"] not in (-1, None): + raise ValueError("NVFP4Config: `group_size` is not used; keep it at `-1`.") + if "weight_block_size" in layer_dict and int(layer_dict["weight_block_size"]) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"NVFP4Config: `weight_block_size` must be `{NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + raw_format = layer_dict.get(FORMAT_FIELD_CODE) + if raw_format is not None: + layer_dict[FORMAT_FIELD_CODE] = _normalize_nvfp4_format(raw_format) + + def quant_linear_init_kwargs(self) -> Dict[str, Any]: + return { + "weight_block_size": self.weight_block_size, + } + + def _update_output_payload(self, out: Dict[str, Any]) -> None: + out[FORMAT_FIELD_CODE] = FORMAT.NVFP4.value + out["weight_block_size"] = self.weight_block_size + + def uses_weight_only_lifecycle(self) -> bool: + return True + @dataclass class BitsAndBytesConfig(PreProcessorConfig): bits: int = field(default=4, metadata={"choices": [4, 8]}) @@ -4078,9 +4223,9 @@ def uses_weight_only_lifecycle(self) -> bool: return True def clone_weight_only_config_for_module( - qcfg: Union[RTNConfig, GGUFConfig, FP8Config, BitsAndBytesConfig], + qcfg: Union[RTNConfig, GGUFConfig, FP8Config, NVFP4Config, BitsAndBytesConfig], module_full_name: str, -) -> Optional[Union[RTNConfig, GGUFConfig, FP8Config, BitsAndBytesConfig]]: +) -> Optional[Union[RTNConfig, GGUFConfig, FP8Config, NVFP4Config, BitsAndBytesConfig]]: if qcfg.dynamic_get(layer_name=module_full_name) is False: return None @@ -4225,11 +4370,12 @@ def _resolve_quantize_config_class(payload: Dict[str, Any]) -> type[BaseQuantize WeightOnlyMethod.RTN, WeightOnlyMethod.GGUF, WeightOnlyMethod.FP8, + WeightOnlyMethod.NVFP4, WeightOnlyMethod.BITSANDBYTES, }: raise ValueError( "QuantizeConfig: unsupported weight-only config. Weight-only export currently supports " - "`rtn`, `gguf`, `fp8`, and `bitsandbytes`." + "`rtn`, `gguf`, `fp8`, `nvfp4`, and `bitsandbytes`." ) if ( format_value == FORMAT.GGUF @@ -4240,6 +4386,8 @@ def _resolve_quantize_config_class(payload: Dict[str, Any]) -> type[BaseQuantize return GGUFConfig if weight_only_method == WeightOnlyMethod.FP8: return FP8Config + if weight_only_method == WeightOnlyMethod.NVFP4: + return NVFP4Config if weight_only_method == WeightOnlyMethod.BITSANDBYTES: return BitsAndBytesConfig if weight_only_method == WeightOnlyMethod.RTN: @@ -4248,6 +4396,12 @@ def _resolve_quantize_config_class(payload: Dict[str, Any]) -> type[BaseQuantize return RTNConfig if method == METHOD.FP8 or format_value == FORMAT.FP8 or _looks_like_fp8_fmt(fp8_storage_fmt): return FP8Config + if method == METHOD.NVFP4: + return NVFP4Config + # `format=nvfp4` with the default/explicit GPTQ method runs the GPTQ algorithm + # on the NVFP4 grid; only `method=nvfp4` selects the direct weight-only path. + if format_value == FORMAT.NVFP4 and method != METHOD.GPTQ: + return NVFP4Config if method == METHOD.BITSANDBYTES or format_value == FORMAT.BITSANDBYTES or _looks_like_bitsandbytes_format(raw_format_value): return BitsAndBytesConfig if method == METHOD.EXL3 or format_value == FORMAT.EXL3: @@ -4276,6 +4430,7 @@ def _known_quantize_config_field_names() -> set[str]: ParoConfig, QQQConfig, FP8Config, + NVFP4Config, BitsAndBytesConfig, EXL3Config, RTNConfig, diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index af04c8d49..1c81954fd 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -21,7 +21,7 @@ from ..looper.named_module import NamedModule from ..quantization import QuantizeConfig -from ..quantization.config import FallbackStrategy, SmoothMSE +from ..quantization.config import FORMAT, FallbackStrategy, SmoothMSE from ..utils.device import get_device from ..utils.logger import setup_logger from ..utils.torch import torch_sync @@ -34,7 +34,7 @@ invert_perm, ) from .npu_linalg import npu_inverse_cholesky_factor -from .quantizer import HF_OPTIMUM, Quantizer +from .quantizer import HF_OPTIMUM, NVFP4Quantizer, Quantizer log = setup_logger() @@ -270,6 +270,8 @@ def validate_module(module): # return any([self.issue_zero_samples, self.issue_nan_hessian, self.issue_non_invertible]) def create_quantizer(self, name: str) -> Quantizer: + if getattr(self.qcfg, "format", None) == FORMAT.NVFP4: + return NVFP4Quantizer(qcfg=self.qcfg, name=name) return Quantizer(qcfg=self.qcfg, name=name) def shape(self): diff --git a/gptqmodel/quantization/quantizer.py b/gptqmodel/quantization/quantizer.py index fa99c5542..f266857d8 100644 --- a/gptqmodel/quantization/quantizer.py +++ b/gptqmodel/quantization/quantizer.py @@ -173,4 +173,52 @@ class QQQQuantizer(Quantizer): def requires_groupwise_processing(self) -> bool: return self.qcfg.group_size == -1 and self.qcfg.sym -__all__ = ["Quantizer"] + +try: + from torchao.prototype.mx_formats.kernels import f4_unpacked_to_f32, f32_to_f4_unpacked +except Exception: + f4_unpacked_to_f32 = None + f32_to_f4_unpacked = None + +_F4_E2M1_MAX = 6.0 +_F8_E4M3_MAX = 448.0 + + +class NVFP4Quantizer(Quantizer): + """Quantizer whose grid is NVFP4: E2M1 values scaled per 16-wide block by E4M3 scales. + + Scale derivation and value rounding mirror torchao's ``nvfp4_quantize`` exactly so + the GPTQ-corrected weight is bit-identical to the packed NVFP4 checkpoint layout. + """ + + def configure(self, *args, **kwargs): + if f32_to_f4_unpacked is None or f4_unpacked_to_f32 is None: + raise RuntimeError("NVFP4Quantizer requires torchao with NVFP4 support.") + super().configure(*args, **kwargs) + + def find_params(self, x, weight=False): + if not weight: + raise NotImplementedError("NVFP4Quantizer only supports weight quantization.") + + dev = x.device + x_flat = x.flatten(1).to(torch.float32) + + # torchao nvfp4_quantize: scale = clamp(amax / 6, e4m3_tiny, 448) cast to E4M3. + amax = x_flat.abs().amax(dim=1) + scale = amax / _F4_E2M1_MAX + scale = torch.clamp(scale, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) + scale = scale.to(torch.float8_e4m3fn).to(torch.float32) + + shape = [-1] + [1] * (len(x.shape) - 1) + self.scale = scale.reshape(shape).to(device=dev) + self.zero = torch.zeros_like(self.scale) + + def quantize(self, x): + orig_dtype = x.dtype + scaled = torch.clamp(x.to(torch.float32) / self.scale, -_F4_E2M1_MAX, _F4_E2M1_MAX) + codes = f32_to_f4_unpacked(scaled.contiguous()) + dequant = f4_unpacked_to_f32(codes) * self.scale + return dequant.to(orig_dtype) + + +__all__ = ["Quantizer", "QQQQuantizer", "NVFP4Quantizer"] diff --git a/gptqmodel/utils/backend.py b/gptqmodel/utils/backend.py index 9d4dde801..301f779e8 100644 --- a/gptqmodel/utils/backend.py +++ b/gptqmodel/utils/backend.py @@ -50,6 +50,9 @@ class BACKEND(str, Enum): # FP8 kernels FP8_TORCH = "fp8_torch" + # NVFP4 kernels + NVFP4_TORCH = "nvfp4_torch" + # GGUF kernels / engines GGUF_TORCH = "gguf_torch" GGUF_TRITON = "gguf_triton" @@ -130,6 +133,9 @@ class PROFILE(str, Enum): "fp8": { BACKEND.TORCH: BACKEND.FP8_TORCH, }, + "nvfp4": { + BACKEND.TORCH: BACKEND.NVFP4_TORCH, + }, "exl3": { BACKEND.EXLLAMA_V3: BACKEND.EXL3_EXLLAMA_V3, BACKEND.TORCH: BACKEND.EXL3_TORCH, diff --git a/gptqmodel/utils/model.py b/gptqmodel/utils/model.py index e58e04f14..88692d74e 100644 --- a/gptqmodel/utils/model.py +++ b/gptqmodel/utils/model.py @@ -40,6 +40,7 @@ DEVICE, EXPERT_INDEX_PLACEHOLDER, SUPPORTS_MODULE_TYPES, + normalize_device, ) from ..nn_modules.qlinear import BaseQuantLinear, GPTQQuantLinear from ..nn_modules.qlinear.exllamav2 import ExllamaV2Linear @@ -622,7 +623,7 @@ def create_quant_module( dtype=dtype, in_features=in_features, out_features=out_features, - device=DEVICE(device) if isinstance(device, str) else device, + device=normalize_device(device) if device is not None else None, adapter=adapter, # TODO FIX ME..need to pass Eora if loaded ) if err is not None: diff --git a/gptqmodel/utils/structure.py b/gptqmodel/utils/structure.py index 6edf2ef33..3839accb7 100644 --- a/gptqmodel/utils/structure.py +++ b/gptqmodel/utils/structure.py @@ -2759,6 +2759,15 @@ def _materialize_direct_meta_tensors( ) -> int: synced = 0 + # Quantized linear tensors are produced by quantization and restored from the + # quant offload directory, never from the source checkpoint. Kernels whose + # buffer names shadow checkpoint tensors (e.g. NVFP4/FP8 `weight`) must not be + # re-aliased from the original dense weights here. + from ..nn_modules.qlinear import BaseQuantLinear + + if isinstance(shell_sub, BaseQuantLinear): + return synced + with torch.inference_mode(): for name, shell_param in dict(shell_sub.named_parameters(recurse=False)).items(): if not _is_meta_tensor(shell_param): diff --git a/scripts/minimax_m3_nvfp4/debug_expert_capture.py b/scripts/minimax_m3_nvfp4/debug_expert_capture.py new file mode 100644 index 000000000..097eec7f3 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/debug_expert_capture.py @@ -0,0 +1,93 @@ +"""Diagnose why layer-3 expert capture saw ~0 tokens during the pilot. + +Loads the model the same way the pilot does, attaches plain forward hooks to all +layer-3 expert Linears, runs ONE 8k calibration sample through the model, and +reports per-expert token counts plus which experts implementation is active. +""" + +import json +import random + +import torch +import torch.nn as nn +from transformers import AutoTokenizer + +MODEL = "/data/sgambhira/models/Minimax-M3-0602" +SNAP = ( + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745" +) + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + offload_to_disk_path="/data/sgambhira/gptqmodel-offload/debug-capture", +) +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("DBG model loaded", flush=True) + +cfg = model.model.config +text_cfg = getattr(cfg, "text_config", cfg) +print("DBG top config _experts_implementation:", getattr(cfg, "_experts_implementation", ""), flush=True) +print("DBG text config _experts_implementation:", getattr(text_cfg, "_experts_implementation", ""), flush=True) + +# Materialize only what a 4-layer forward needs (weights are lazy/meta otherwise), +# using the same API the looper uses. +lm = model.model.model.language_model +device = torch.device("cpu") +for sub in [lm.embed_tokens, lm.rotary_emb, lm.norm, *lm.layers[:4]]: + model.shell_module_materialize(sub, device) +print("DBG materialized embed/rotary/norm/layers0-3", flush=True) +# truncate the stack so the forward stops after layer 3 +lm.layers = nn.ModuleList(list(lm.layers[:4])) + +layer3 = lm.layers[3] +experts_mod = layer3.mlp.experts +print("DBG experts module class:", type(experts_mod).__name__, flush=True) +print("DBG experts forward qualname:", type(experts_mod).forward.__qualname__, flush=True) +has_3d = any(p.ndim == 3 for p in experts_mod.parameters(recurse=False)) +print("DBG experts has 3D params:", has_3d, flush=True) + +counts = {} + +def make_hook(name): + def hook(mod, inp, out): + counts[name] = counts.get(name, 0) + inp[0].shape[0] + return hook + +hooks = [] +for e in range(128): + child = getattr(experts_mod, str(e), None) + if child is None: + continue + for proj in ("gate_proj", "up_proj", "down_proj"): + m = getattr(child, proj, None) + if isinstance(m, nn.Linear): + hooks.append(m.register_forward_hook(make_hook(f"{e}.{proj}"))) +print("DBG hooks attached:", len(hooks), flush=True) + +# one 8k calibration sample +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(l) for l in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +r = random.Random(42).sample(rows, 64)[0] +messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, +] +text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) +ids = tokenizer(text, add_special_tokens=False, return_tensors="pt")["input_ids"][:, :8192] +print("DBG sample tokens:", ids.shape[1], flush=True) + +with torch.inference_mode(): + lm(input_ids=ids) + +gate_counts = {k: v for k, v in counts.items() if k.endswith("gate_proj")} +total = sum(gate_counts.values()) +nonzero = len(gate_counts) +print(f"DBG experts hit: {nonzero}/128 | total routed tokens (gate): {total:,} | expected ~{ids.shape[1]*4:,}", flush=True) +top = sorted(gate_counts.items(), key=lambda kv: -kv[1])[:8] +print("DBG top experts:", top, flush=True) +print("DBG DONE", flush=True) diff --git a/scripts/minimax_m3_nvfp4/p03_load_dryrun.py b/scripts/minimax_m3_nvfp4/p03_load_dryrun.py new file mode 100644 index 000000000..b30ba5e52 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/p03_load_dryrun.py @@ -0,0 +1,116 @@ +"""Phase 0.3: GPTQModel.load dry run for Minimax-M3-0602 NVFP4 plan. + +Verifies: (a) load routes to the quantize path and MXFP8 decodes to bf16, +(b) runtime module naming (block_sparse_moe/w1 vs mlp/gate_proj), (c) the +experts-only dynamic rules select the expected 16,512 modules. + +Run: .venv/bin/python scripts/minimax_m3_nvfp4/p03_load_dryrun.py +See: docs/minimax_m3_nvfp4_quantization_plan.md +""" + +import json + +import torch +import torch.nn as nn + +MODEL = "/data/sgambhira/models/Minimax-M3-0602" +SKIP_LAYERS = "22|23|24|25|26|27|28|32|33|34|35|36|37|38" + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + dynamic={ + r"-:.*self_attn\..*": {}, + r"-:.*shared_experts\..*": {}, + r"-:.*layers\.(0|1|2)\.mlp\..*": {}, + rf"-:.*layers\.({SKIP_LAYERS})\..*": {}, + }, + moe_vram_strategy="balanced", +) +print("P03 qcfg class:", type(qcfg).__name__, "format:", qcfg.format, "device:", qcfg.device, flush=True) + +# Non-quantized models always load CPU-side (lazy); GPUs are used per-layer +# during quantization. Do not pass device= when qcfg.device is set/auto. +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("P03 loaded model class:", type(model.model).__name__, flush=True) + +# (b) runtime module names for layer 3 +names_l3 = [ + (n, type(m).__name__) + for n, m in model.model.named_modules() + if ".layers.3." in n and isinstance(m, nn.Linear) +] +print("P03 layer-3 linear count:", len(names_l3), flush=True) +for n, t in names_l3[:12]: + print("P03 l3 module:", n, t, flush=True) +expert_markers = { + "mlp.experts": sum(1 for n, _ in names_l3 if "mlp.experts." in n), + "block_sparse_moe.experts": sum(1 for n, _ in names_l3 if "block_sparse_moe.experts." in n), +} +print("P03 expert naming counts (layer 3):", expert_markers, flush=True) + +# (a) dequant spot check on one expert projection +expert_name = next((n for n, _ in names_l3 if ".experts.0." in n), None) +print("P03 spot-check module:", expert_name, flush=True) +mod = dict(model.model.named_modules())[expert_name] +w = mod.weight +print("P03 module weight dtype/shape/device:", w.dtype, tuple(w.shape), w.device, flush=True) + +if w.device.type != "meta": + from safetensors import safe_open + + index = json.load(open(f"{MODEL}/model.safetensors.index.json"))["weight_map"] + cands = [ + k + for k in index + if ".layers.3." in k + and ".experts.0." in k + and k.endswith(".weight") + and "scale" not in k + and "shared" not in k + ] + print("P03 checkpoint candidates:", cands[:3], flush=True) + matched = False + for ck in cands: + with safe_open(f"{MODEL}/{index[ck]}", framework="pt", device="cpu") as f: + fp8 = f.get_tensor(ck) + scale_u8 = f.get_tensor(ck.replace(".weight", ".weight_scale_inv")) + if tuple(fp8.shape) != tuple(w.shape): + continue + scale = torch.pow(2.0, scale_u8.to(torch.float32) - 127.0) + manual = fp8.to(torch.float32) * scale.repeat_interleave(32, dim=1) + diff = (w.detach().to(torch.float32).cpu() - manual).abs().max().item() + print(f"P03 dequant max abs diff vs manual mxfp8 ({ck}):", diff, flush=True) + matched = True + break + if not matched: + print("P03 no shape-matching checkpoint candidate found", flush=True) + +# (c) dynamic selection counts over all language-model layers +included, excluded = [], [] +for n, m in model.model.named_modules(): + if not isinstance(m, nn.Linear) or ".layers." not in n: + continue + if "visual" in n or "vision" in n: + continue + if qcfg.dynamic_get(layer_name=n) is False: + excluded.append(n) + else: + included.append(n) +experts_included = [n for n in included if ".experts." in n and "shared" not in n] +non_expert_included = [n for n in included if n not in set(experts_included)] +print("P03 total linears:", len(included) + len(excluded), flush=True) +print("P03 included expert linears:", len(experts_included), "(expect 16512)", flush=True) +print("P03 included NON-expert linears:", len(non_expert_included), flush=True) +for n in non_expert_included[:10]: + print("P03 non-expert included (should be empty or router-only):", n, flush=True) +skip_hits = [ + n + for n in included + if any(f".layers.{x}." in n for x in "22 23 24 25 26 27 28 32 33 34 35 36 37 38".split()) +] +print("P03 included modules in skip layers (expect 0):", len(skip_hits), flush=True) +print("P03 DONE", flush=True) diff --git a/scripts/minimax_m3_nvfp4/p1_pilot_layer3.py b/scripts/minimax_m3_nvfp4/p1_pilot_layer3.py new file mode 100644 index 000000000..1087e9abf --- /dev/null +++ b/scripts/minimax_m3_nvfp4/p1_pilot_layer3.py @@ -0,0 +1,123 @@ +"""Phase 1: single-layer (layer 3) GPTQ->NVFP4 pilot for Minimax-M3-0602. + +Quantizes ONLY layer-3 routed experts (384 modules) with the full 64-trajectory +HLE calibration set. Gates checked here before the full run: + - per-expert calibration coverage / RTN-fallback rate under native routing + - loss sanity (also validates the MXFP8->bf16 decode numerically) + - saved tensor layout (uint8 packed weight + float8_e4m3fn weight_scale) + - what non-quantized modules are written as (plan open-decision #1) + +Run: .venv/bin/python scripts/minimax_m3_nvfp4/p1_pilot_layer3.py +See: docs/minimax_m3_nvfp4_quantization_plan.md +""" + +import json +import random + +import torch +from transformers import AutoTokenizer + +MODEL = "/data/sgambhira/models/Minimax-M3-0602" +OUT = "/data/sgambhira/models/Minimax-M3-0602-NVFP4-pilot-l3" +SNAP = ( + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745" +) +SEED = 42 +NUM_SAMPLES = 64 +# 80GB H100 headroom (user-confirmed): keep the FIRST 8k tokens per sample. The sdpa +# sparse-attention path materializes dense per-head masks (64 x n^2 x 2B) plus a bool +# intermediate — ~9GB + ~4GB at 8k vs 32GB + 17GB at 16k, which OOMed. +MAX_TOKENS_PER_SAMPLE = 8192 + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + dynamic={ + r"-:.*self_attn\..*": {}, + r"-:.*shared_experts\..*": {}, + r"-:.*layers\.(0|1|2)\.mlp\..*": {}, + # pilot: skip every layer except 3 + r"-:.*layers\.([4-9]|[1-5][0-9])\..*": {}, + }, + moe_vram_strategy="balanced", + # /tmp is on the ~full 124G root disk; module offload must go to /data (15T free). + offload_to_disk_path="/data/sgambhira/gptqmodel-offload/pilot-l3", +) +print("P1 qcfg:", type(qcfg).__name__, qcfg.format, flush=True) + +# --- calibration: 64 seed-42 HLE trajectories, pre-tokenized (matches P0.2) --- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(l) for l in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) + +calibration = [] +total_tokens = 0 +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + enc = tokenizer(text, add_special_tokens=False, return_tensors="pt") + input_ids = enc["input_ids"][:, :MAX_TOKENS_PER_SAMPLE] + attention_mask = enc["attention_mask"][:, :MAX_TOKENS_PER_SAMPLE] + calibration.append({"input_ids": input_ids, "attention_mask": attention_mask}) + total_tokens += input_ids.shape[1] +truncated = sum(1 for c in calibration if c["input_ids"].shape[1] == MAX_TOKENS_PER_SAMPLE) +print( + f"P1 calibration: {len(calibration)} samples, {total_tokens:,} tokens " + f"({truncated} truncated to first {MAX_TOKENS_PER_SAMPLE})", + flush=True, +) + +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("P1 model loaded", flush=True) + +model.quantize(calibration, batch_size=1, backend=BACKEND.TORCH) +print("P1 quantize complete", flush=True) + +model.save(OUT) +print("P1 saved to", OUT, flush=True) + +# --- post-checks on the saved checkpoint --- +from safetensors import safe_open # noqa: E402 + +index = json.load(open(f"{OUT}/model.safetensors.index.json"))["weight_map"] + +def _tensor(name): + with safe_open(f"{OUT}/{index[name]}", framework="pt", device="cpu") as f: + return f.get_tensor(name) + +l3_keys = sorted(k for k in index if ".layers.3.mlp.experts.0." in k) +print("P1 saved layer-3 expert-0 tensors:", l3_keys, flush=True) +for k in l3_keys: + t = _tensor(k) + print(f"P1 saved {k}: {t.dtype} {tuple(t.shape)}", flush=True) + +# non-quantized passthrough: layer-4 expert + layer-3 attention +for probe in [ + next((k for k in index if ".layers.4.mlp.experts.0." in k and k.endswith(".weight")), None), + next((k for k in index if ".layers.3.self_attn.q_proj." in k and k.endswith(".weight")), None), +]: + if probe: + t = _tensor(probe) + print(f"P1 passthrough {probe}: {t.dtype} {tuple(t.shape)}", flush=True) + +# dequant round-trip vs torchao for one quantized expert +wk = next((k for k in index if ".layers.3.mlp.experts.0.down_proj.weight" == k.split("model.", 1)[-1] or k.endswith("layers.3.mlp.experts.0.down_proj.weight")), None) +sk = wk.replace(".weight", ".weight_scale") if wk else None +if wk and sk in index: + from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor + + w, s = _tensor(wk), _tensor(sk) + nv = NVFP4Tensor(w, s, block_size=16, orig_dtype=torch.bfloat16) + deq = nv.dequantize(torch.bfloat16) if not hasattr(nv, "to_dtype") else nv.to_dtype(torch.bfloat16) + print(f"P1 nvfp4 round-trip ok: {wk} -> dequant {tuple(deq.shape)} finite={torch.isfinite(deq).all().item()}", flush=True) +else: + print("P1 WARNING: could not locate quantized down_proj weight/scale pair", flush=True) + +print("P1 DONE", flush=True) diff --git a/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py b/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py new file mode 100644 index 000000000..28aee9c19 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py @@ -0,0 +1,167 @@ +"""Phase 1b: per-expert routing coverage across all MoE layers of Minimax-M3-0602. + +Streams the model layer-window by layer-window through one GPU (materialize -> +forward all samples -> free), capturing each MoE layer's routed expert +assignments (top_k_index) at the experts-module boundary. Native top-4 routing, +64 seed-42 HLE samples truncated to their first 8k tokens — identical to the +Phase 1 pilot calibration. + +Output: JSON matrix [57 layers x 128 experts] of routed token counts plus a +per-layer concentration summary. +""" + +import json +import random + +import torch +import torch.nn as nn +from transformers import AutoTokenizer + +MODEL = "/data/sgambhira/models/Minimax-M3-0602" +SNAP = ( + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745" +) +OUT_JSON = "/data/sgambhira/GPTQModel/scripts/minimax_m3_nvfp4/coverage_layer_expert.json" +SEED = 42 +NUM_SAMPLES = 64 +MAX_TOKENS = 8192 +WINDOW = 5 +NUM_EXPERTS = 128 +# GPU 0 is occupied by the Phase-1 pilot process during its save; use an idle GPU. +DEVICE = torch.device("cuda:4") + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + offload_to_disk_path="/data/sgambhira/gptqmodel-offload/coverage", +) +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("COV model loaded", flush=True) + +lm = model.model.model.language_model +num_layers = len(lm.layers) +all_layers = list(lm.layers) + +# --- calibration inputs (identical to pilot) --- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(l) for l in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) +inputs = [] +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + ids = tokenizer(text, add_special_tokens=False, return_tensors="pt")["input_ids"][:, :MAX_TOKENS] + inputs.append(ids) +total_tokens = sum(i.shape[1] for i in inputs) +print(f"COV calibration: {len(inputs)} samples, {total_tokens:,} tokens", flush=True) + +# --- routed-assignment capture at the experts-module boundary --- +counts = {} # layer_idx -> LongTensor[NUM_EXPERTS] + +def make_experts_hook(layer_idx): + def pre_hook(module, args, kwargs): + idx = None + if len(args) >= 2 and torch.is_tensor(args[1]): + idx = args[1] + elif "top_k_index" in kwargs: + idx = kwargs["top_k_index"] + if idx is None: + return + c = torch.bincount(idx.reshape(-1).to("cpu", torch.int64), minlength=NUM_EXPERTS) + counts[layer_idx] = counts.get(layer_idx, torch.zeros(NUM_EXPERTS, dtype=torch.long)) + c + return pre_hook + +hook_handles = [] +moe_layers = [] +for li, layer in enumerate(all_layers): + experts = getattr(getattr(layer, "mlp", None), "experts", None) + if isinstance(experts, nn.Module): + moe_layers.append(li) + hook_handles.append(experts.register_forward_pre_hook(make_experts_hook(li), with_kwargs=True)) +print(f"COV hooked {len(moe_layers)} MoE layers: {moe_layers[:4]}...{moe_layers[-2:]}", flush=True) + +# --- windowed streaming forward --- +# hidden states enter each window via inputs_embeds; rotary/causal-mask are +# position-derived so per-window model forwards are exact. +embed = model.shell_module_materialize(lm.embed_tokens, DEVICE) +model.shell_module_materialize(lm.rotary_emb, DEVICE) +model.shell_module_materialize(lm.norm, DEVICE) + +hidden = [] +with torch.inference_mode(): + for ids in inputs: + hidden.append(embed(ids.to(DEVICE)).to("cpu")) +print("COV embedded all samples", flush=True) + +orig_layers = lm.layers + +for w_start in range(0, num_layers, WINDOW): + w_layers = all_layers[w_start : w_start + WINDOW] + # materialization resolves module paths against the model tree, so the full + # layer list must be attached while materializing; truncate only for forward. + lm.layers = orig_layers + for layer in w_layers: + model.shell_module_materialize(layer, DEVICE) + lm.layers = nn.ModuleList(w_layers) + + # capture the last window layer's output as the next window's input + window_out = {} + def tail_hook(module, args, output): + window_out["h"] = (output[0] if isinstance(output, tuple) else output).detach() + tail = w_layers[-1].register_forward_hook(tail_hook) + + with torch.inference_mode(): + for si in range(len(hidden)): + h = hidden[si].to(DEVICE) + pos = torch.arange(h.shape[1], device=DEVICE).unsqueeze(0) + lm(inputs_embeds=h, position_ids=pos, use_cache=False) + hidden[si] = window_out["h"].to("cpu") + + tail.remove() + for layer in w_layers: + layer.to_empty(device="meta") + torch.cuda.empty_cache() + print(f"COV window done: layers {w_start}-{w_start + len(w_layers) - 1}", flush=True) + +lm.layers = orig_layers + +for h in hook_handles: + h.remove() + +# --- summarize --- +matrix = {str(li): counts.get(li, torch.zeros(NUM_EXPERTS, dtype=torch.long)).tolist() for li in moe_layers} +payload = { + "model": MODEL, + "samples": NUM_SAMPLES, + "max_tokens": MAX_TOKENS, + "seed": SEED, + "total_tokens": total_tokens, + "top_k": 4, + "counts": matrix, +} +with open(OUT_JSON, "w") as f: + json.dump(payload, f) +print(f"COV wrote {OUT_JSON}", flush=True) + +threshold = 0.005 * total_tokens +print("COV layer | experts>0 | experts>=fallback_thr | top1_share | top8_share", flush=True) +for li in moe_layers: + c = torch.tensor(matrix[str(li)], dtype=torch.float64) + tot = c.sum().item() + if tot == 0: + print(f"COV {li:>3} | ZERO CAPTURE", flush=True) + continue + srt = c.sort(descending=True).values + print( + f"COV {li:>3} | {(c > 0).sum().item():>3} | {(c >= threshold).sum().item():>3} | " + f"{srt[0].item() / tot:5.1%} | {srt[:8].sum().item() / tot:5.1%}", + flush=True, + ) +print("COV DONE", flush=True) diff --git a/tests/test_nvfp4_export.py b/tests/test_nvfp4_export.py new file mode 100644 index 000000000..72c9ba820 --- /dev/null +++ b/tests/test_nvfp4_export.py @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +"""NVFP4 export coverage: GPTQ-on-NVFP4-grid and weight-only NVFP4 packing. + +1. `NVFP4Quantizer` must be bit-identical to torchao's ``nvfp4_quantize`` grid. +2. `TorchNVFP4Linear` packing must be lossless against the scales produced by + the GPTQ loop (error feedback mutates columns after group scales are fixed, + so export must reuse the loop's scales, never re-derive them). +3. End-to-end: GPTQ quantize -> save -> reload -> forward on a tiny model, + with the checkpoint stored in the NVFP4 layout (packed uint8 weight + + float8_e4m3fn `weight_scale`). +""" + +import os +from pathlib import Path + + +# Keep this test on CPU so it works on CPU-only runners. +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + +import pytest +import torch +import torch.nn as nn +from safetensors import safe_open +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.trainers import WordLevelTrainer +from transformers import PreTrainedTokenizerFast +from transformers.models.llama.modeling_llama import LlamaConfig, LlamaForCausalLM + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig +from gptqmodel.nn_modules.qlinear.fp4 import TorchNVFP4Linear, quantize_nvfp4_weight +from gptqmodel.quantization.config import FORMAT, METHOD, GPTQConfig, NVFP4Config +from gptqmodel.quantization.quantizer import NVFP4Quantizer + + +try: + from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor, nvfp4_quantize +except Exception: + NVFP4Tensor = None + nvfp4_quantize = None + +pytestmark = [ + pytest.mark.cpu, + pytest.mark.skipif(NVFP4Tensor is None, reason="torchao NVFP4 support required"), +] + +BLOCK = 16 + + +def _nvfp4_dequant(packed: torch.Tensor, scales: torch.Tensor, dtype=torch.float32) -> torch.Tensor: + nv = NVFP4Tensor(packed, scales, block_size=BLOCK, orig_dtype=dtype) + to_dtype = getattr(nv, "to_dtype", None) + if callable(to_dtype): + return to_dtype(dtype) + return nv.dequantize(dtype) + + +# --------------------------------------------------------------------------- +# Config resolution +# --------------------------------------------------------------------------- + + +def test_format_nvfp4_resolves_to_gptq_config(): + cfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + assert isinstance(cfg, GPTQConfig) + assert cfg.method == METHOD.GPTQ + assert cfg.format == FORMAT.NVFP4 + assert cfg.sym is True + assert cfg.desc_act is False + assert cfg.act_group_aware is False + + payload = cfg.to_dict() + reloaded = QuantizeConfig.from_quant_config(payload) + assert isinstance(reloaded, GPTQConfig) + assert reloaded.format == FORMAT.NVFP4 + + +def test_method_nvfp4_resolves_to_weight_only_config(): + cfg = QuantizeConfig(quant_method="nvfp4") + assert isinstance(cfg, NVFP4Config) + assert cfg.uses_weight_only_lifecycle() is True + assert cfg.bits == 4 + assert cfg.weight_block_size == BLOCK + + reloaded = QuantizeConfig.from_quant_config(cfg.to_dict()) + assert isinstance(reloaded, NVFP4Config) + + shorthand = QuantizeConfig(weight_only="nvfp4") + assert isinstance(shorthand, NVFP4Config) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"bits": 8, "group_size": 16}, + {"bits": 4, "group_size": 128}, + {"bits": 4, "group_size": 16, "sym": False}, + {"bits": 4, "group_size": 16, "desc_act": True}, + {"bits": 4, "group_size": 16, "act_group_aware": True}, + ], +) +def test_format_nvfp4_rejects_incompatible_gptq_settings(kwargs): + with pytest.raises(ValueError): + QuantizeConfig(format="nvfp4", **kwargs) + + +# --------------------------------------------------------------------------- +# Quantizer grid exactness +# --------------------------------------------------------------------------- + + +def test_nvfp4_quantizer_matches_torchao_grid(): + torch.manual_seed(0) + W = torch.randn(64, 128, dtype=torch.float32) + + scales_ref, packed_ref = nvfp4_quantize(W, block_size=BLOCK) + deq_ref = _nvfp4_dequant(packed_ref, scales_ref) + + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + out = torch.empty_like(W) + scales = [] + for g in range(0, W.shape[1], BLOCK): + block = W[:, g : g + BLOCK] + quantizer.find_params(block, weight=True) + out[:, g : g + BLOCK] = quantizer.quantize(block) + scales.append(quantizer.scale) + scales = torch.cat(scales, dim=1) + + assert torch.equal(out, deq_ref) + assert torch.equal( + scales.to(torch.float8_e4m3fn).view(torch.uint8), + scales_ref.view(torch.uint8), + ) + + +# --------------------------------------------------------------------------- +# Kernel packing +# --------------------------------------------------------------------------- + + +def test_direct_pack_matches_torchao(): + torch.manual_seed(1) + lin = nn.Linear(128, 64, bias=True) + + module = TorchNVFP4Linear( + bits=4, + group_size=-1, + sym=True, + desc_act=False, + in_features=128, + out_features=64, + bias=True, + register_buffers=True, + ) + module.pack_original(linear=lin, scales=None, zeros=None) + + scales_ref, packed_ref = nvfp4_quantize(lin.weight.detach().float(), block_size=BLOCK) + assert torch.equal(module.weight, packed_ref.view(torch.uint8)) + assert torch.equal(module.weight_scale.view(torch.uint8), scales_ref.view(torch.uint8)) + + deq = module.dequantize_weight(dtype=torch.float32).T + assert torch.allclose(deq, _nvfp4_dequant(packed_ref, scales_ref), atol=1e-3) + + +def test_gptq_scale_pack_is_lossless(): + """Error feedback shrinks/grows columns after the group scale is fixed; packing + against the loop's scales must reproduce the GPTQ output weight exactly.""" + + torch.manual_seed(2) + IN, OUT = 128, 64 + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + W = torch.randn(OUT, IN, dtype=torch.float32) + Q = torch.empty_like(W) + scales = [] + for g in range(0, IN, BLOCK): + block = W[:, g : g + BLOCK] + quantizer.find_params(block, weight=True) + # simulate GPTQ error feedback mutating the block after scale is fixed + block = block + 0.05 * torch.randn_like(block) + Q[:, g : g + BLOCK] = quantizer.quantize(block) + scales.append(quantizer.scale) + scales = torch.cat(scales, dim=1) + + lin = nn.Linear(IN, OUT, bias=False) + lin.weight.data = Q.clone() + module = TorchNVFP4Linear( + bits=4, + group_size=16, + sym=True, + desc_act=False, + in_features=IN, + out_features=OUT, + bias=False, + register_buffers=True, + ) + module.pack_original(linear=lin, scales=scales, zeros=None) + + deq = module.dequantize_weight(dtype=torch.float32).T + assert torch.equal(deq, Q) + + +def test_quantize_nvfp4_weight_rejects_bad_shapes(): + with pytest.raises(ValueError): + quantize_nvfp4_weight(torch.randn(4, 24)) # 24 % 16 != 0 + with pytest.raises(ValueError): + quantize_nvfp4_weight(torch.randn(4, 16, 2)) + + +# --------------------------------------------------------------------------- +# End-to-end: GPTQ quantize -> save -> reload -> forward +# --------------------------------------------------------------------------- + +_CALIBRATION_TEXTS = [ + "tiny nvfp4 calibration sample one with enough tokens to survive minimum length filtering", + "tiny nvfp4 calibration sample two exercising the gptq loop over the fp4 grid cleanly", + "another synthetic calibration example that is intentionally verbose so filtering keeps it", +] * 2 + + +def _build_local_tokenizer(model_dir: Path) -> PreTrainedTokenizerFast: + tokenizer = Tokenizer(WordLevel(unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + trainer = WordLevelTrainer(special_tokens=["[PAD]", "[UNK]", "[BOS]", "[EOS]"]) + tokenizer.train_from_iterator(_CALIBRATION_TEXTS, trainer=trainer) + + fast_tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + bos_token="[BOS]", + eos_token="[EOS]", + unk_token="[UNK]", + pad_token="[PAD]", + ) + fast_tokenizer.save_pretrained(model_dir) + return fast_tokenizer + + +def _build_tiny_llama_fixture(model_dir: Path) -> tuple[LlamaConfig, PreTrainedTokenizerFast]: + config = LlamaConfig( + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=4, + vocab_size=128, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + ) + model = LlamaForCausalLM(config) + model.save_pretrained(model_dir) + tokenizer = _build_local_tokenizer(model_dir) + return config, tokenizer + + +def _build_calibration_dataset(tokenizer: PreTrainedTokenizerFast) -> list[dict[str, object]]: + dataset = [] + for text in _CALIBRATION_TEXTS: + encoded = tokenizer(text, return_tensors="pt") + dataset.append( + { + "input_ids": encoded["input_ids"], + "attention_mask": encoded["attention_mask"], + } + ) + return dataset + + +@pytest.mark.slow +def test_gptq_nvfp4_end_to_end(tmp_path: Path): + model_dir = tmp_path / "native" + quantized_dir = tmp_path / "quantized" + + config, tokenizer = _build_tiny_llama_fixture(model_dir) + calibration = _build_calibration_dataset(tokenizer) + + quantize_config = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + device="cpu", + ) + assert isinstance(quantize_config, GPTQConfig) + + model = GPTQModel.load( + str(model_dir), + quantize_config=quantize_config, + backend=BACKEND.TORCH, + ) + model.quantize( + calibration, + batch_size=1, + backend=BACKEND.TORCH, + calibration_data_min_length=1, + ) + model.save(quantized_dir) + + # Checkpoint layout: packed uint8 fp4 weight + float8_e4m3fn weight_scale. + shards = sorted(quantized_dir.glob("*.safetensors")) + assert shards, "no safetensors shard written" + weight_key = "model.layers.0.self_attn.q_proj.weight" + scale_key = "model.layers.0.self_attn.q_proj.weight_scale" + with safe_open(shards[0], framework="pt", device="cpu") as f: + keys = set(f.keys()) + assert weight_key in keys and scale_key in keys, sorted(keys) + weight = f.get_tensor(weight_key) + scale = f.get_tensor(scale_key) + hidden = config.hidden_size + assert weight.dtype is torch.uint8 + assert weight.shape == (hidden, hidden // 2) + assert scale.dtype is torch.float8_e4m3fn + assert scale.shape == (hidden, hidden // BLOCK) + + # The saved tensors must decode with the torchao NVFP4 layout. + _ = _nvfp4_dequant(weight, scale) + + reloaded = GPTQModel.load(str(quantized_dir), device="cpu") + modules = dict(reloaded.named_modules()) + for suffix in ("q_proj", "k_proj", "v_proj", "o_proj"): + name = f"model.model.layers.0.self_attn.{suffix}" + assert isinstance(modules[name], TorchNVFP4Linear), name + for suffix in ("gate_proj", "up_proj", "down_proj"): + name = f"model.model.layers.0.mlp.{suffix}" + assert isinstance(modules[name], TorchNVFP4Linear), name + + encoded = tokenizer("nvfp4 forward smoke", return_tensors="pt") + with torch.inference_mode(): + out = reloaded.model(input_ids=encoded["input_ids"]) + assert out.logits.shape[-1] == config.vocab_size + assert torch.isfinite(out.logits).all() + + +@pytest.mark.slow +def test_weight_only_nvfp4_end_to_end(tmp_path: Path): + """RTN-style direct NVFP4 export (`method=nvfp4`) without calibration.""" + + model_dir = tmp_path / "native" + quantized_dir = tmp_path / "quantized" + + config, tokenizer = _build_tiny_llama_fixture(model_dir) + + quantize_config = QuantizeConfig(quant_method="nvfp4", device="cpu") + assert isinstance(quantize_config, NVFP4Config) + + model = GPTQModel.load( + str(model_dir), + quantize_config=quantize_config, + backend=BACKEND.TORCH, + ) + model.quantize(calibration=None, backend=BACKEND.TORCH) + model.save(quantized_dir) + + reloaded = GPTQModel.load(str(quantized_dir), device="cpu") + modules = dict(reloaded.named_modules()) + name = "model.model.layers.0.mlp.down_proj" + assert isinstance(modules[name], TorchNVFP4Linear), name + + encoded = tokenizer("nvfp4 weight only smoke", return_tensors="pt") + with torch.inference_mode(): + out = reloaded.model(input_ids=encoded["input_ids"]) + assert torch.isfinite(out.logits).all() From 3e01e1c8101516e9e7bd1754327a632c233c5876 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 16:21:43 -0700 Subject: [PATCH 02/14] fix(loader): apply mxfp8 block scales when materializing fp8 checkpoints LazyTurtle copied checkpoint tensors into shell params with a plain dtype cast. For an mxfp8 checkpoint the stored tensor is only the F8_E4M3 mantissa payload; the magnitude lives in a sibling weight_scale_inv holding E8M0 exponent bytes, one per 32-element block. Casting to bf16 silently dropped that scale, so every quantized weight was wrong by 2^(e-127). The scale was never even requested: t_params comes from the shell module's named_parameters(), and a plain nn.Linear has no weight_scale_inv attribute, so _resolve_checkpoint_tensor_source was never asked for it. Measured on Minimax-M3-0602, layers.3.self_attn.q_proj: the materialized tensor was bit-identical to the raw fp8 payload (absmax 448, std 148.6) instead of the scaled weight (absmax 0.2812, std 0.0381). Downstream, MoE routing collapsed onto 6.8 of 128 experts per layer with one expert taking every token in 55 of 57 layers; with scales applied it reaches 127.9 of 128 with no expert above a 17.5% share. dequantize_fp8() could not be reused: _expand_scale treats scale_inv as a float multiplier and broadcasts by rank, so E8M0 bytes would multiply by 120.0 rather than 2^-7, and a [8192,192] scale will not broadcast against a [8192,6144] weight. Adds dequantize_block_fp8 / block_scale_multiplier, which decode E8M0 and expand each dim by target_dim // scale_dim. Verified bit-for-bit (maxdiff 0) against modelopt's MXFP8QTensor.dequantize on attention, MoE expert, and shared-expert weights. Routes every materialization read through _read_checkpoint_tensor, covering both _copy_checkpoint_tensors_into_submodule and the parallel _materialize_direct_meta_tensors path (param/buffer, single and concat). Layers in ignored_layers carry no scale tensor, so the lookup misses and they keep the plain-cast path -- which is why the F32 routers loaded correctly even while the bug was live. Not covered: _load_checkpoint_tensors_for_module_path still returns raw payloads. It has 8 callers in models/base.py and returns the _scale_inv entries alongside the weights, so dequantizing there could double-apply; those call sites need review separately. Co-Authored-By: Claude Opus 5 (1M context) --- gptqmodel/quantization/dtype.py | 66 ++++++++++++++++++++++++++++ gptqmodel/utils/structure.py | 77 ++++++++++++++++++++++++++++++--- 2 files changed, 137 insertions(+), 6 deletions(-) diff --git a/gptqmodel/quantization/dtype.py b/gptqmodel/quantization/dtype.py index 15e34f7a7..1627e256b 100644 --- a/gptqmodel/quantization/dtype.py +++ b/gptqmodel/quantization/dtype.py @@ -42,6 +42,9 @@ "dequantize_fp8", "dequantize_f8_e4m3", "dequantize_f4_e2m1", + "dequantize_block_fp8", + "block_scale_multiplier", + "fp8_storage_dtypes", "is_fp4_packed_dtype", ] @@ -681,6 +684,69 @@ def dequantize_fp8( ) +# E8M0 is a scale format, never a weight payload, so it is excluded here: a +# tensor stored as e8m0 is metadata and must not be dequantized as if it were +# quantized data. +_FP8_STORAGE_DTYPE_NAMES = ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") +_FP8_STORAGE_DTYPES = tuple( + getattr(torch, name) for name in _FP8_STORAGE_DTYPE_NAMES if hasattr(torch, name) +) + + +def fp8_storage_dtypes() -> tuple[torch.dtype, ...]: + """FP8 dtypes that can carry quantized weight payloads (excludes E8M0 scales).""" + + return _FP8_STORAGE_DTYPES + + +def block_scale_multiplier(scale: torch.Tensor, target_shape: torch.Size) -> torch.Tensor: + """Expand a block scale so it multiplies elementwise against ``target_shape``. + + ``uint8`` scales are MX E8M0 exponents, whose multiplier is ``2**(e - 127)``. + Float scales are already multipliers (block-fp8 checkpoints store them that + way). Each dim is repeated by ``target_dim // scale_dim``, so ``[1, 32]`` + and ``[128, 128]`` block layouts are both handled without special-casing. + """ + + if scale.dtype == torch.uint8: + mult = torch.exp2(scale.to(torch.float32) - 127.0) + else: + mult = scale.to(torch.float32) + + while mult.ndim < len(target_shape): + mult = mult.unsqueeze(-1) + if mult.ndim != len(target_shape): + raise ValueError( + f"block scale rank {mult.ndim} exceeds weight rank {len(target_shape)}" + ) + + for dim, (want, have) in enumerate(zip(target_shape, mult.shape)): + if have == want: + continue + if have == 0 or want % have != 0: + raise ValueError( + f"block scale dim {dim} ({have}) does not evenly divide weight dim ({want})" + ) + mult = mult.repeat_interleave(want // have, dim=dim) + return mult + + +def dequantize_block_fp8( + tensor: torch.Tensor, + scale: torch.Tensor, + *, + target_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize a block-scaled FP8 weight to ``target_dtype``. + + Unlike :func:`dequantize_fp8`, this applies a *block* scale (one entry per + fixed-size group of elements) and decodes E8M0 exponent bytes, which is what + mxfp8 checkpoints store alongside each weight. + """ + + return (tensor.to(torch.float32) * block_scale_multiplier(scale, tensor.shape)).to(target_dtype) + + def dequantize_f4_e2m1( tensor: torch.Tensor, *, diff --git a/gptqmodel/utils/structure.py b/gptqmodel/utils/structure.py index 3839accb7..bd4d3077c 100644 --- a/gptqmodel/utils/structure.py +++ b/gptqmodel/utils/structure.py @@ -44,6 +44,7 @@ from safetensors import safe_open from torch import nn +from ..quantization.dtype import dequantize_block_fp8, fp8_storage_dtypes from ..utils.logger import setup_logger @@ -2294,6 +2295,58 @@ def _load_checkpoint_tensors_for_module_path( tensors[rel_name] = handler.get_tensor(full_name) return tensors + def _materialize_dtype(self) -> torch.dtype: + """Compute dtype for dequantized checkpoint tensors.""" + + for cfg in (getattr(self.config, "text_config", None), self.config): + dtype = getattr(cfg, "torch_dtype", None) or getattr(cfg, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str) and hasattr(torch, dtype): + resolved = getattr(torch, dtype) + if isinstance(resolved, torch.dtype): + return resolved + return torch.bfloat16 + + def _scale_tensor_name(self, full_name: str) -> Optional[str]: + """Checkpoint name of the block scale paired with ``full_name``, if any.""" + + candidate = f"{full_name}_scale_inv" + return candidate if candidate in self._weight_map else None + + def _read_checkpoint_tensor(self, handler, shard_path: str, full_name: str) -> torch.Tensor: + """Read a checkpoint tensor, applying block scales to FP8 payloads. + + An FP8 checkpoint stores only the mantissa payload; the magnitude lives + in a sibling ``*_scale_inv`` tensor. Copying the payload into a bf16 + parameter with a plain dtype cast silently drops that scale, so the + dequant has to happen here, before the tensor reaches the target param. + + Layers the exporter left unquantized (``ignored_layers``) carry no scale + tensor, so the lookup misses and they keep the plain-cast path. + """ + + tensor = handler.get_tensor(full_name) + if tensor.dtype not in fp8_storage_dtypes(): + return tensor + + scale_name = self._scale_tensor_name(full_name) + if scale_name is None: + return tensor + + scale_shard = self._weight_map.get(scale_name) + if scale_shard is None: + return tensor + + scale_path = os.path.join(self.model_local_path, scale_shard) + if scale_path == shard_path: + scale = handler.get_tensor(scale_name) + else: + with safe_open(scale_path, framework="pt", device="cpu") as scale_handler: + scale = scale_handler.get_tensor(scale_name) + + return dequantize_block_fp8(tensor, scale, target_dtype=self._materialize_dtype()) + def _copy_checkpoint_tensors_into_submodule( self, *, @@ -2414,7 +2467,9 @@ def _copy_checkpoint_tensors_into_submodule( ) shard_path = os.path.join(self.model_local_path, shard) with safe_open(shard_path, framework="pt", device="cpu") as handler: - parts.append(handler.get_tensor(full_name)) + parts.append( + self._read_checkpoint_tensor(handler, shard_path, full_name) + ) try: tensor = torch.cat(parts, dim=concat_dim).contiguous() @@ -2511,7 +2566,9 @@ def _copy_checkpoint_tensors_into_submodule( rel_name=rel_name, modules_by_name=modules_by_name, ) - checkpoint_tensor = handler.get_tensor(full_name) + checkpoint_tensor = self._read_checkpoint_tensor( + handler, shard_path, full_name + ) tensor = self._transform_checkpoint_tensor( checkpoint_tensor, expert_index=expert_index, @@ -2791,7 +2848,9 @@ def _materialize_direct_meta_tensors( )) source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - parts.append(handler.get_tensor(full_name)) + parts.append( + self._read_checkpoint_tensor(handler, source_path, full_name) + ) try: source_param = torch.cat(parts, dim=concat_dim).contiguous() @@ -2853,7 +2912,9 @@ def _materialize_direct_meta_tensors( source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - checkpoint_param = handler.get_tensor(full_name) + checkpoint_param = self._read_checkpoint_tensor( + handler, source_path, full_name + ) source_param = self._transform_checkpoint_tensor( checkpoint_param, expert_index=expert_index, @@ -2928,7 +2989,9 @@ def _materialize_direct_meta_tensors( )) source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - parts.append(handler.get_tensor(full_name)) + parts.append( + self._read_checkpoint_tensor(handler, source_path, full_name) + ) try: source_buffer = torch.cat(parts, dim=concat_dim).contiguous() @@ -2988,7 +3051,9 @@ def _materialize_direct_meta_tensors( source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - checkpoint_buffer = handler.get_tensor(full_name) + checkpoint_buffer = self._read_checkpoint_tensor( + handler, source_path, full_name + ) source_buffer = self._transform_checkpoint_tensor( checkpoint_buffer, expert_index=expert_index, From 96cc70621b86477ba01c10d678944f9796507f6c Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 16:47:44 -0700 Subject: [PATCH 03/14] feat(gptq): accumulate the Hessian on quantized activations for w4a4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GPTQ minimizes the layer output error. For weight-only quantization inference computes Ŵ·X, so H = 2·X·Xᵀ is the right curvature. Under w4a4 the kernel computes Ŵ·Q(X), and a Hessian built from X solves for inputs that never occur. At 4-bit activations that gap is large; at 16-bit it does not exist, which is why this has not mattered until now. Activations reach the solve through exactly one path -- process_batch -> compute_hessian_xtx -> self.H -- so quantizing there also propagates to dead-column detection, act-order, act-group-aware grouping, and the damped Cholesky inverse. Adds GPTQConfig.act_format / act_group_size. act_format=None keeps today's weight-only behavior exactly: no activation quantizer is constructed and process_batch is unchanged. The quantizer is applied explicitly rather than captured by a hook. This is deliberate: a forward_pre_hook fires before a module's forward body and so observes activations *before* any input quantizer, which would silently accumulate a clean Hessian while every config flag looked correct. The regression test asserts H actually differs (~4.8% relative on random input) precisely to catch that failure mode. Ordering in process_batch matters twice: the quantizer runs before the TP zero-padding, so pad columns never enter a 16-wide block's amax, and after any input-scale transform already baked into `inp`, since that is what the kernel quantizes. fake_quantize_nvfp4 wraps torchao's nvfp4_quantize + NVFP4Tensor.dequantize rather than reimplementing the grid, so it cannot drift from the weight path. Following that path, per_tensor_scale is omitted -- block scales only. Scope: this is tier A of docs/w4a4_hessian.md, min ‖(W−Ŵ)·Q(X)‖², matching Model-Optimizer's GPTQHelper. It is an approximation: the exact objective min ‖W·X − Ŵ·Q(X)‖² carries a W·(X−Q(X)) bias no Ŵ can cancel, which needs a second accumulator (tier B, cf. Quark GPTAQ's ∆(XXᵀ)). Known gap: NVFP4 activations use a static per-tensor global scale that is calibrated, plus dynamic per-block scales. Nothing populates the global scale yet -- Model-Optimizer gets it from a separate max_calibrate stage that GPTQModel has no equivalent of. Until that lands, the calibration grid is close to but not identical to a static-global-scale serving kernel. Co-Authored-By: Claude Opus 5 (1M context) --- docs/w4a4_hessian.md | 233 +++++++++++++++++++++++++++++++ gptqmodel/quantization/config.py | 31 ++++ gptqmodel/quantization/dtype.py | 40 ++++++ gptqmodel/quantization/gptq.py | 33 ++++- tests/test_w4a4_hessian.py | 137 ++++++++++++++++++ 5 files changed, 473 insertions(+), 1 deletion(-) create mode 100644 docs/w4a4_hessian.md create mode 100644 tests/test_w4a4_hessian.py diff --git a/docs/w4a4_hessian.md b/docs/w4a4_hessian.md new file mode 100644 index 000000000..9c4b80a5a --- /dev/null +++ b/docs/w4a4_hessian.md @@ -0,0 +1,233 @@ +# w4a4: accumulating the GPTQ Hessian on quantized activations + +## The problem + +GPTQ minimizes the layer output error. For weight-only quantization inference +computes `Ŵ·X`, so the objective and its Hessian are: + +``` +min_Ŵ ‖ W·X − Ŵ·X ‖² = min_Ŵ ‖ (W − Ŵ)·X ‖² -> H = 2·X·Xᵀ +``` + +Under w4a4 inference computes `Ŵ·Q(X)`. If calibration accumulates `H` on `X` +while the kernel feeds `Q(X)`, the solved weights are optimal for inputs the +model never sees. At 4-bit activations that gap is large — unlike w4a16, where +`Q(X) = X` and the question does not arise. + +In GPTQModel, activations enter in exactly one place: `GPTQ.process_batch` +(`gptqmodel/quantization/gptq.py`) reshapes the layer input and hands it to +`compute_hessian_xtx`; `add_batch` accumulates into `self.H`. After that `X` is +never referenced again. `self.H` feeds dead-column detection, act-order +permutation, act-group-aware grouping, and the damped Cholesky inverse — so +substituting at accumulation time propagates to all of them. + +## Two tiers + +| tier | objective | statistics | reference implementation | +|---|---|---|---| +| **A** | `min ‖(W−Ŵ)·Q(X)‖²` | `H` only | Model-Optimizer `GPTQHelper` | +| **B** | `min ‖W·X − Ŵ·Q(X)‖²` | `H` + `∆(XXᵀ)` | Quark GPTAQ / Qronos | + +Tier A substitutes the input and treats activation-quantization error as an +irreducible constant. It is an approximation: the true residual `W·X − Ŵ·Q(X)` +is **not** of the form `(W − Ŵ)·M` for any `M`, because the bias `W·(X − Q(X))` +cannot be cancelled by any choice of `Ŵ`. It is nonetheless strictly better than +the status quo. + +Tier B is the exact objective and absorbs that bias, at the cost of a second +accumulator and changes to the `fasterquant` solve. + +Tier A is a prerequisite for B. Recommendation: land A, measure against a w4a4 +perplexity baseline, then decide on B. + +--- + +## What the reference implementations do + +### Model-Optimizer — does exactly tier A + +`modelopt/torch/quantization/utils/calib_utils.py`, `GPTQHelper.setup()`: + +```python +def hessian_forward(self, input, *args, **kwargs): + inp = input.to_local() if hasattr(input, "to_local") else input + if self.input_quantizer is not None and self.input_quantizer.is_enabled: + hessian_input = self.input_quantizer(inp) # Q(X) + else: + hessian_input = inp # X + gptq_helper.hessian, gptq_helper.n_samples = update_hessian( + hessian_input, gptq_helper.hessian, gptq_helper.n_samples + ) + return self._forward_no_gptq_hessian(input, *args, **kwargs) + +bind_forward_method(self.module, hessian_forward, self.CACHE_NAME) +``` + +Three things to take from this: + +1. **One branch, no separate w4a4 path.** The Hessian goes on `Q(X)` whenever the + input quantizer is enabled. +2. **It reuses the same `input_quantizer` object the inference path uses** + (`conversion.py` does `input = self.input_quantizer(input)`), so the + calibration scheme cannot drift from the serving scheme. +3. **It patches `forward`, it does not use a hook.** This is essential. A + `forward_pre_hook` fires before the forward body and therefore sees the input + *before* the internal `input_quantizer` call. **The quantized activation + cannot be captured with a hook — the quantizer must be applied explicitly.** + A port that uses hooks will silently accumulate clean activations. + +`update_hessian` itself is standard (`H += (2/n)·XXᵀ` with incremental +rescaling), including a zero-token guard for MoE experts that receive no tokens. + +#### How the static activation scale is resolved + +NVFP4 activations use a **static per-tensor global scale** plus **dynamic +per-16-block E4M3 scales**. The global scale is a property of the whole +calibration set, which appears to create a chicken-and-egg with the calibration +pass itself. + +Model-Optimizer resolves it by **ordering, not by a two-pass GPTQ**: the input +quantizer's amax is calibrated in an earlier, separate stage (`max_calibrate`), +so by the time the GPTQ stage runs the quantizer is already calibrated and +enabled. GPTQ stays single-pass. This works because the scale belongs to the +quantizer, not to GPTQ. + +#### Paths in Model-Optimizer that deliberately do *not* do this + +- `local_hessian_calibrate` (`model_calib.py`) feeds **clean** activations. Its + hooks are `register_forward_pre_hook`s capturing `args[0]` — pre-quantizer, per + the point above. Correct there, because it is a different object: a per-cin-block + Hessian used to weight a *weight*-error metric + (`local_hessian_error(x, xq)` computes `dw = x − xq` over `cout`). +- **AWQ-lite / AWQ-clip** *do* use quantized activations, gated on + `is_input_quantized = module.input_quantizer.is_enabled`, then searching + against `self.input_quantizer(input)`. Note AWQ-clip calls `max_calibrate` on + the current batch and uses the result immediately — a running max. Acceptable + for a scale search; not advisable under a matrix that gets Cholesky-inverted. + +### Quark — a different axis, plus the tier-B machinery + +Quark's GPTQ (`quark/torch/algorithm/gptq/gptq.py`) exposes +`add_batch_quantized` / `add_batch_nonquantized`, which resembles this +distinction but is not it. In Quark, "quantized" means **the input produced by +running the block with already-quantized preceding weights** — sequential +weight-error propagation. Evidence: + +- The hook is a `register_forward_hook` capturing `input[0]`, i.e. pre-quantizer. +- The two passes differ by `layer_inputs` vs `orig_layer_inputs`, and the clean + pass runs under `RestoreOriginalWeights(layer)`. Restoring *weights* is what + produces the clean input, so the axis is weights, not activations. +- Plain GPTQ's `add_batch_nonquantized` raises + `ValueError("We don't need it for GPTQ")`. + +So Quark is not a precedent for activation-quantized Hessians. It does, however, +implement the asymmetric objective needed for tier B, in +`quark/torch/algorithm/gptaq/gptaq.py` (GPTAQ, arXiv 2504.02692) and +`qronos/qronos.py`: + +```python +# add_batch_quantized: H = X̃ X̃ᵀ +self.H += inp.matmul(inp.t()) + +# add_batch_nonquantized: ∆(XXᵀ) = (X − X̃) X̃ᵀ +delta_input = original_input - self.q_input +self.delta_X_Xt += delta_input.matmul(self.q_input.t()) +``` + +That is `min_Ŵ ‖W·X − Ŵ·X̃‖²`, structurally identical to the w4a4 objective with +`X̃ = Q(X)`; `∆(XXᵀ)` is the cross term tier A drops. GPTAQ also runs two +`block_forward`s per batch, so a second pass is accepted practice. + +--- + +## Proposed change to GPTQModel (tier A) + +### Config + +`grep` for `act_bits|activation_bits|act_quant|activation_scheme` in +`gptqmodel/quantization/config.py` returns nothing — GPTQModel has no notion of +activation quantization today. This adds a capability, not a flag. Add to +`GPTQConfig`: + +```python +act_format: Optional[FORMAT] = None # None = weight-only (today's behavior) +act_group_size: int = 16 # NVFP4 block +``` + +Validation: `act_format=nvfp4` should require `format=nvfp4`, and should reject +`mock_quantization` for the same reason the weight path does. + +### `gptq.py` + +Insert in `process_batch`, after `reshaped_inp = reshaped_inp.contiguous()` and +**before** the `_tp_pad_cols` zero-padding: + +```python +reshaped_inp = reshaped_inp.contiguous() + +# w4a4: the kernel sees Q(X), so the Hessian must describe Q(X), not X. +# Quantize before TP zero-padding so pad columns do not enter block statistics. +if self._act_quantizer is not None: + reshaped_inp = self._act_quantizer(reshaped_inp) + +if self._tp_pad_cols: + ... +``` + +Ordering matters twice: + +- **Before TP padding** — `_tp_pad_cols` appends zero columns; quantizing after + would let those zeros enter a 16-wide block's amax and shift its scale. +- **After any input-scale / smoothing transform** — `Q` must see the + post-transform activation, because that is what the kernel quantizes. + +### The activation quantizer + +`NVFP4Quantizer` (`gptqmodel/quantization/quantizer.py`) is **not** reusable: +`find_params` raises `NotImplementedError` for `weight=False`, and it derives one +scale per row of a group-sliced weight. Activations need per-token × per-16-block +along the last dim of a `(tokens, cin)` tensor. + +Use torchao's `nvfp4_quantize(x, block_size=16)` — already imported in +`gptqmodel/quantization/dtype.py` and used by the weight path — then dequantize +back to the staging dtype so `compute_hessian_xtx` is untouched. + +Following Model-Optimizer, the quantizer should be an object shared with (or +provably identical to) whatever the serving kernel applies, rather than a +reimplementation kept in sync by hand. + +## Consequences to watch + +- **Dead columns.** `dead = diag(H) == 0` currently sets `H[dead, dead] = 1`, + silently skipping correction for that column. Small-magnitude channels can + flush to zero under FP4, creating dead columns that did not exist in the clean + Hessian. Log a count when activation quantization is enabled. +- **act-order / act-group-aware.** Both sort on `diag(H)`. With `Q(X)` the + ordering is by quantized-activation importance — arguably more correct for + w4a4, but group assignments will differ from a weight-only run, so results are + not comparable across the flag. +- **Cost.** One quantize+dequantize per batch, negligible against the `XᵀX` GEMM, + but it adds a full-size fp32 temporary — relevant given the existing OOM + fallback in `process_batch`. +- **MoE.** Each expert sees only its routed tokens. With few tokens per expert a + dynamic scale is noisier; the static global scale mitigates this, which is + another argument for matching the serving scheme exactly. + +## Open questions + +1. Exact `amax → global_scale` formula for activations, and whether the amax is a + plain max or percentile/MSE-clipped. Must match the serving kernel or every + block scale is off by a constant. +2. Whether the input-scale (`inscale1`) transform applies to activations at + inference. If so `Q` sits after it, and any amax must be measured + post-transform. +3. Whether `X̃` should compose both error sources — upstream weight quantization + *and* activation quantization, i.e. `X̃ = Q(X_weightpath)`. That is what the + deployed model sees, and neither reference implementation does both at once. + +## Verification plan + +- A layer where `H_quant` and `H_clean` differ measurably (otherwise the change + is a no-op and something is misconfigured — e.g. quantizing via a hook). +- End-to-end: quantizing with `H_quant` should beat `H_clean` on w4a4 perplexity. + Without this the change is unfalsifiable. diff --git a/gptqmodel/quantization/config.py b/gptqmodel/quantization/config.py index 32c18b238..978da7f70 100644 --- a/gptqmodel/quantization/config.py +++ b/gptqmodel/quantization/config.py @@ -3253,6 +3253,18 @@ class GPTQConfig(PreProcessorConfig): metadata={"help": "Skip heavy computations for fast model loading validation"}, ) hessian: Optional[HessianConfig] = field(default_factory=HessianConfig) + act_format: Optional[FORMAT] = field( + default=None, + metadata={ + "help": "Activation format the deployed kernel will use (w4a4). When set, the " + "GPTQ Hessian is accumulated on Q(X) instead of X, so the solved weights " + "are optimal for the inputs the kernel actually sees. None = weight-only." + }, + ) + act_group_size: int = field( + default=NVFP4_WEIGHT_BLOCK_SIZE, + metadata={"help": "Activation quantization block size, along the last dim."}, + ) def allowed_quant_methods(self) -> Tuple[METHOD, ...]: return (METHOD.GPTQ,) @@ -3310,6 +3322,25 @@ def __post_init__(self): if float(self.mse or 0.0) > 0.0: raise ValueError("QuantizeConfig: `format=nvfp4` does not support `mse` scale search.") + if self.act_format is not None: + # w4a4: the Hessian is built from Q(X), so the activation grid has to be one + # we can actually reproduce here. + if self.act_format != FORMAT.NVFP4: + raise ValueError( + f"QuantizeConfig: `act_format={self.act_format}` is not supported; " + "only `act_format=nvfp4` is implemented." + ) + if self.format != FORMAT.NVFP4: + raise ValueError( + "QuantizeConfig: `act_format=nvfp4` requires `format=nvfp4`." + ) + if self.act_group_size <= 0: + raise ValueError("QuantizeConfig: `act_group_size` must be positive.") + if self.mock_quantization: + raise ValueError( + "QuantizeConfig: `act_format` does not support `mock_quantization=True`." + ) + def _resolve_activation_ordering( self, desc_act_user_value: Optional[bool], diff --git a/gptqmodel/quantization/dtype.py b/gptqmodel/quantization/dtype.py index 1627e256b..4dec74a57 100644 --- a/gptqmodel/quantization/dtype.py +++ b/gptqmodel/quantization/dtype.py @@ -44,6 +44,7 @@ "dequantize_f4_e2m1", "dequantize_block_fp8", "block_scale_multiplier", + "fake_quantize_nvfp4", "fp8_storage_dtypes", "is_fp4_packed_dtype", ] @@ -699,6 +700,45 @@ def fp8_storage_dtypes() -> tuple[torch.dtype, ...]: return _FP8_STORAGE_DTYPES +def fake_quantize_nvfp4( + x: torch.Tensor, + *, + block_size: int = 16, + per_tensor_scale: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Quantize to the NVFP4 grid and immediately dequantize, preserving shape/dtype. + + Used to make calibration statistics describe the activations a w4a4 kernel + actually consumes. Blocks run along the last dim, matching the weight path in + ``quantize_to_nvfp4``; ``per_tensor_scale`` is the optional static global + scale (omitted, as in the weight path, when the grid has block scales only). + """ + + if NVFP4Tensor is None or nvfp4_quantize is None: + raise RuntimeError("fake_quantize_nvfp4 requires torchao with NVFP4 support.") + + if x.shape[-1] % block_size != 0: + raise ValueError( + f"NVFP4 activation quantization needs the last dim ({x.shape[-1]}) " + f"to be a multiple of block_size ({block_size})." + ) + + orig_dtype = x.dtype + scales, packed = nvfp4_quantize( + x.to(torch.float32).contiguous(), + block_size=block_size, + per_tensor_scale=per_tensor_scale, + ) + tensor = NVFP4Tensor( + packed, + scales, + block_size, + orig_dtype, + per_tensor_scale=per_tensor_scale, + ) + return tensor.dequantize(orig_dtype) + + def block_scale_multiplier(scale: torch.Tensor, target_shape: torch.Size) -> torch.Tensor: """Expand a block scale so it multiplies elementwise against ``target_shape``. diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index 1c81954fd..14d9aeb7d 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -11,7 +11,7 @@ import sys import threading import time -from typing import Dict, Optional, Tuple +from typing import Callable, Dict, Optional, Tuple import numpy as np import torch @@ -25,6 +25,7 @@ from ..utils.device import get_device from ..utils.logger import setup_logger from ..utils.torch import torch_sync +from .dtype import fake_quantize_nvfp4 from .fallback_smooth import mse_optimal_quant, smooth_block from .gar import ( compose_final_perm, @@ -215,6 +216,7 @@ def __init__(self, module: nn.Module, qcfg: Optional[QuantizeConfig] = None): self.nsamples = 0 self.quantizer = self.create_quantizer(name=self.name) + self._act_quantizer = self._create_act_quantizer() # fwd counter self.fwd_counter = 0 @@ -269,6 +271,28 @@ def validate_module(module): # def has_hessian_issues(self) -> bool: # return any([self.issue_zero_samples, self.issue_nan_hessian, self.issue_non_invertible]) + def _create_act_quantizer(self) -> Optional[Callable[[torch.Tensor], torch.Tensor]]: + """Build the activation fake-quantizer for w4a4, or ``None`` for weight-only. + + Under w4a4 the kernel consumes ``Q(X)``, so the Hessian must be built from + ``Q(X)`` rather than ``X`` — otherwise GPTQ solves for inputs that never + occur. Returning ``None`` preserves the weight-only behaviour exactly. + """ + + act_format = getattr(self.qcfg, "act_format", None) + if act_format is None: + return None + + if act_format != FORMAT.NVFP4: + raise ValueError(f"GPTQ: unsupported `act_format={act_format}`.") + + block_size = int(getattr(self.qcfg, "act_group_size", 16)) + + def _quantize_act(x: torch.Tensor) -> torch.Tensor: + return fake_quantize_nvfp4(x, block_size=block_size) + + return _quantize_act + def create_quantizer(self, name: str) -> Quantizer: if getattr(self.qcfg, "format", None) == FORMAT.NVFP4: return NVFP4Quantizer(qcfg=self.qcfg, name=name) @@ -530,6 +554,13 @@ def process_batch(self, inp: torch.Tensor) -> Tuple[int, Optional[torch.Tensor], # Delay dtype conversion until we materialize Hessian chunks to avoid unnecessary temporaries reshaped_inp = reshaped_inp.contiguous() + + # w4a4: the kernel consumes Q(X), so the Hessian must describe Q(X), not X. + # Applied before the TP zero-padding below so the pad columns never enter a + # block's amax, and after any input-scale transform baked into `inp`. + if self._act_quantizer is not None: + reshaped_inp = self._act_quantizer(reshaped_inp) + if self._tp_pad_cols: pad = reshaped_inp.new_zeros((reshaped_inp.shape[0], self._tp_pad_cols)) reshaped_inp = torch.cat((reshaped_inp, pad), dim=1) diff --git a/tests/test_w4a4_hessian.py b/tests/test_w4a4_hessian.py new file mode 100644 index 000000000..37e50e65c --- /dev/null +++ b/tests/test_w4a4_hessian.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +"""w4a4: the GPTQ Hessian must be accumulated on quantized activations. + +Under w4a4 the kernel consumes ``Q(X)``, so a Hessian built from ``X`` solves for +inputs that never occur. The load-bearing test here is +``test_hessian_differs_from_clean_activations``: the natural way to get this +wrong is to capture activations with a ``forward_pre_hook``, which fires before +any input quantizer and therefore silently accumulates clean activations. That +bug leaves every config flag looking correct, so it is only detectable by +asserting the Hessian actually changed. + +See `docs/w4a4_hessian.md`. +""" + +import os + + +# Keep this test on CPU so it works on CPU-only runners. +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + +import pytest +import torch +import torch.nn as nn + +from gptqmodel.quantization import QuantizeConfig +from gptqmodel.quantization.config import FORMAT +from gptqmodel.quantization.dtype import fake_quantize_nvfp4 +from gptqmodel.quantization.gptq import GPTQ + + +IN_FEATURES = 64 +OUT_FEATURES = 32 +BLOCK = 16 + + +def _qcfg(act: bool) -> QuantizeConfig: + kwargs = dict(bits=4, group_size=BLOCK, format="nvfp4", sym=True, desc_act=False) + if act: + kwargs.update(act_format=FORMAT.NVFP4, act_group_size=BLOCK) + return QuantizeConfig(**kwargs) + + +def _hessian_for(act: bool, inp: torch.Tensor) -> torch.Tensor: + torch.manual_seed(0) + module = nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False) + gptq = GPTQ(module=module, qcfg=_qcfg(act)) + gptq.add_batch(inp, torch.empty(0)) + return gptq.finalize_hessian().clone() + + +@pytest.fixture(scope="module") +def calib_input() -> torch.Tensor: + torch.manual_seed(1234) + return torch.randn(4, 128, IN_FEATURES) + + +def test_fake_quantize_nvfp4_round_trip(): + """Quantize+dequantize preserves layout but lands values on the FP4 grid.""" + + torch.manual_seed(0) + x = torch.randn(8, IN_FEATURES) + q = fake_quantize_nvfp4(x, block_size=BLOCK) + + assert q.shape == x.shape + assert q.dtype == x.dtype + assert torch.isfinite(q).all() + assert not torch.equal(q, x), "fake quantization was a no-op" + # E2M1 has 16 codes, so one block can never hold more distinct values. + assert len(torch.unique(q[0, :BLOCK])) <= 16 + + +def test_fake_quantize_nvfp4_rejects_ragged_last_dim(): + """A last dim that is not block-divisible must raise, not silently mis-block.""" + + with pytest.raises(ValueError, match="multiple of block_size"): + fake_quantize_nvfp4(torch.randn(4, BLOCK + 1), block_size=BLOCK) + + +def test_weight_only_config_leaves_activations_untouched(): + """Without `act_format`, no activation quantizer exists at all.""" + + gptq = GPTQ(module=nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False), qcfg=_qcfg(act=False)) + assert gptq._act_quantizer is None + + +def test_act_format_installs_quantizer(): + gptq = GPTQ(module=nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False), qcfg=_qcfg(act=True)) + assert gptq._act_quantizer is not None + + +def test_hessian_differs_from_clean_activations(calib_input): + """The regression that matters: `act_format` must actually change the Hessian. + + A hook-based implementation would see pre-quantizer activations and produce a + Hessian identical to the weight-only one while appearing configured. + """ + + h_clean = _hessian_for(act=False, inp=calib_input) + h_quant = _hessian_for(act=True, inp=calib_input) + + assert h_clean.shape == h_quant.shape + assert not torch.allclose(h_clean, h_quant), ( + "Hessian is identical with and without activation quantization — " + "activations are not being quantized before accumulation" + ) + rel = ((h_clean - h_quant).norm() / h_clean.norm()).item() + assert rel > 1e-3, f"Hessian changed by only {rel:.2%}; expected a visible FP4 effect" + + +def test_quantized_hessian_is_well_formed(calib_input): + """Quantizing activations must not break the properties the solve relies on.""" + + h = _hessian_for(act=True, inp=calib_input) + + assert torch.isfinite(h).all() + assert torch.allclose(h, h.T, atol=1e-4), "Hessian lost symmetry" + assert (torch.diag(h) > 0).all(), "Hessian gained a dead column under activation quantization" + + +def test_hessian_is_deterministic(calib_input): + """Same inputs must give the same Hessian — the solve is not robust to drift.""" + + assert torch.equal(_hessian_for(act=True, inp=calib_input), _hessian_for(act=True, inp=calib_input)) + + +@pytest.mark.parametrize( + "kwargs, match", + [ + (dict(bits=4, group_size=128, format="gptq", act_format=FORMAT.NVFP4), "requires `format=nvfp4`"), + (dict(bits=4, group_size=BLOCK, format="nvfp4", act_format=FORMAT.NVFP4, act_group_size=0), "act_group_size"), + ], +) +def test_config_rejects_unsupported_activation_settings(kwargs, match): + with pytest.raises(ValueError, match=match): + QuantizeConfig(**kwargs) From 16422dd1803f08e60af9657cb8d1250d67618400 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 16:55:23 -0700 Subject: [PATCH 04/14] feat(gptq): load NVFP4's static activation scale from an amax sidecar NVFP4 activations are scaled twice: a static per-tensor global scale frozen at calibration, and dynamic per-16-block E4M3 scales. Only the block scales can be derived while accumulating the Hessian, so w4a4 calibration was quantizing against a block-only grid -- close to, but not the same grid a static-global kernel consumes. The global scale cannot be computed here by construction: it is a property of the whole calibration set, and we are inside the pass that would produce it. Both reference pipelines resolve this by ordering rather than by a two-pass GPTQ -- Model-Optimizer freezes it in an earlier max_calibrate stage, and tore-quant emits it from Phase 1 (`ptq.py --emit-amax-sidecar`) for Phase 2 to consume. This follows the latter: read the sidecar the pipeline already produces instead of adding a calibration stage. Scale formula mirrors tore-quant's amax_to_scale exactly: global_scale = amax / 6.0 / 448.0 # amax / F4_E2M1_MAX / F8_E4M3_MAX The amax is a plain max in that pipeline, not percentile- or MSE-clipped. act_amax_key_rules exists because the scale is keyed by *input site*, not by module. w1 and w3 consume the same MoE input and therefore share one entry (`layer{N}.moe_input`) while w2 sees the post-SwiGLU activation and gets its own; q/k/v share their attention input by the same logic. A per-module lookup would hand siblings different scales than the kernel applies -- the exact mismatch the sidecar exists to prevent -- so a test pins the sharing. Resolution is lazy rather than done in __init__: `self.name` is not final at construction time (the looper names modules afterwards), so an early lookup reads the wrong key. This was caught by the new tests, and in production would have surfaced as a wrong scale rather than a crash. A missing key raises instead of falling back to block-only scales; a silent fallback would quantize against a different grid than the kernel with no signal. Non-finite and non-positive amax are rejected. The sidecar loader is lru_cached so thousands of modules read the file once, and accepts both a flat {key: amax} mapping and the pipeline's {"amax": {...}} shape. Co-Authored-By: Claude Opus 5 (1M context) --- docs/w4a4_hessian.md | 63 +++++++++++++++-- gptqmodel/quantization/config.py | 24 +++++++ gptqmodel/quantization/gptq.py | 80 ++++++++++++++++++++- tests/test_w4a4_hessian.py | 115 +++++++++++++++++++++++++++++++ 4 files changed, 276 insertions(+), 6 deletions(-) diff --git a/docs/w4a4_hessian.md b/docs/w4a4_hessian.md index 9c4b80a5a..692e30562 100644 --- a/docs/w4a4_hessian.md +++ b/docs/w4a4_hessian.md @@ -213,15 +213,68 @@ reimplementation kept in sync by hand. dynamic scale is noisier; the static global scale mitigates this, which is another argument for matching the serving scheme exactly. +## The static global scale: how `tore-quant` does it + +`tore-quant` runs the production NVFP4 pipeline and answers the scale questions +concretely. + +**Formula** (`src/tore_quant/convert_mxfp4_to_nvfp4.py`): + +```python +def amax_to_scale(amax: torch.Tensor) -> torch.Tensor: + """NVFP4 second-level scale formula.""" + return amax.float() / 6.0 / 448.0 +``` + +i.e. `global_scale = amax / (F4_E2M1_MAX × F8_E4M3_MAX)`, serialized per module +as a fp32 scalar `.input_scale`. + +**The amax is a plain max, not clipped.** It comes from modelopt max-calibration +buffers, distilled by `save_amax_sidecar` (`src/tore_quant/ptq.py:525`). + +**Two phases, not one pass.** Phase 1 (`ptq.py --emit-amax-sidecar`) writes +`amax_per_layer.json`; Phase 2 (`convert_mxfp4_to_nvfp4`) consumes it. Same shape +as Model-Optimizer: the global scale is frozen by an earlier stage, so the +quantization step never has to solve the chicken-and-egg. + +**The scale is shared by input site, not per module** — the part that matters +most for GPTQModel: + +``` +layer{L}.moe_input <- max input amax over routed-expert w1/w3 +layer{L}.w2_input <- max input amax over routed-expert w2 +``` + +`w1` and `w3` consume the same MoE input, so they share one scale; `w2` sees the +post-SwiGLU activation and gets its own. A peer-max sync +(`fixup_moe_expert_amax`) then makes the value identical across every expert in a +layer and across ranks, so a per-layer scalar loses nothing. + +By the same logic `q/k/v` would share their attention input. + +**MoE calibration caveat, from `--calib-all-experts` (default False):** forcing +every token through every expert "distorts per-expert input amax". Native top-k +routing plus the peer-max sync is the supported path. + +### Consequence for this implementation + +Tier A currently passes `per_tensor_scale=None`, so it uses block scales only. +Closing that gap needs more than a number: `GPTQ` is per-module and has no way to +express "these modules share an activation scale". A faithful implementation +needs the scale keyed by **input site** and supplied from a sidecar, either + +- loaded from an `amax_per_layer.json`-style file (matches the existing + pipeline, no new calibration stage), or +- collected by a GPTQModel-side amax pass with sibling-sharing logic. + +The first is much cheaper and reuses a file the pipeline already produces. + ## Open questions -1. Exact `amax → global_scale` formula for activations, and whether the amax is a - plain max or percentile/MSE-clipped. Must match the serving kernel or every - block scale is off by a constant. -2. Whether the input-scale (`inscale1`) transform applies to activations at +1. Whether the input-scale (`inscale1`) transform applies to activations at inference. If so `Q` sits after it, and any amax must be measured post-transform. -3. Whether `X̃` should compose both error sources — upstream weight quantization +2. Whether `X̃` should compose both error sources — upstream weight quantization *and* activation quantization, i.e. `X̃ = Q(X_weightpath)`. That is what the deployed model sees, and neither reference implementation does both at once. diff --git a/gptqmodel/quantization/config.py b/gptqmodel/quantization/config.py index 978da7f70..288d552ed 100644 --- a/gptqmodel/quantization/config.py +++ b/gptqmodel/quantization/config.py @@ -3277,6 +3277,24 @@ def default_desc_act(self) -> bool: def __post_init__(self): desc_act_user_value = self.desc_act + act_amax_path: Optional[str] = field( + default=None, + metadata={ + "help": "Path to an activation amax sidecar (JSON) used to derive NVFP4's static " + "per-tensor global scale as amax/6/448. Accepts a flat {key: amax} mapping " + "or {'amax': {key: amax}}. Without it only block scales are applied, which " + "will not match a kernel that uses a static global activation scale." + }, + ) + act_amax_key_rules: Optional[List[Tuple[str, str]]] = field( + default=None, + metadata={ + "help": "Ordered (regex, replacement) pairs mapping a module name to its sidecar " + "key, first match wins; default is to look up the module name itself. " + "Modules sharing an input must map to the same key -- e.g. w1/w3 both to " + "`layer{N}.moe_input` -- because they share one activation scale." + }, + ) act_group_aware_user_value = self.act_group_aware super().__post_init__() @@ -3372,6 +3390,12 @@ def _update_meta_payload(self, meta_payload: Dict[str, Any]) -> None: "alpha": self.gptaq.alpha, "device": device if isinstance(device, str) else str(device), } + elif self.act_amax_path is not None or self.act_amax_key_rules is not None: + # A sidecar with no activation quantization silently does nothing; that is + # far more likely a misconfiguration than an intent. + raise ValueError( + "QuantizeConfig: `act_amax_path`/`act_amax_key_rules` require `act_format` to be set." + ) else: device = self.foem.device meta_payload["foem"] = { diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index 14d9aeb7d..97e984cc7 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -6,8 +6,11 @@ # Based on original gptq algorithm and code from https://github.com/IST-DASLab/gptq import contextlib +import functools +import json import math import os +import re import sys import threading import time @@ -130,6 +133,28 @@ def _device_supports_bfloat16(device: torch.device) -> bool: return support +@functools.lru_cache(maxsize=8) +def _load_act_amax_table(path: str) -> Dict[str, float]: + """Load an activation amax sidecar, cached so thousands of modules read it once. + + Accepts a flat ``{key: amax}`` mapping or ``{"amax": {key: amax}}``, the shape + the NVFP4 conversion pipeline emits. + """ + + with open(path, "r") as handle: + payload = json.load(handle) + + if isinstance(payload, dict) and isinstance(payload.get("amax"), dict): + payload = payload["amax"] + + if not isinstance(payload, dict): + raise ValueError( + f"GPTQ: activation amax sidecar '{path}' must be a mapping, or contain an 'amax' mapping." + ) + + return {str(k): float(v) for k, v in payload.items()} + + def get_number_of_rows_and_cols(layer: nn.Module): # return layer.weight.shape[0], np.prod(layer.weight.shape[1:]) if isinstance(layer, NamedModule): @@ -271,6 +296,50 @@ def validate_module(module): # def has_hessian_issues(self) -> bool: # return any([self.issue_zero_samples, self.issue_nan_hessian, self.issue_non_invertible]) + def _resolve_act_global_scale(self) -> Optional[torch.Tensor]: + """NVFP4's static per-tensor activation scale, from the amax sidecar. + + NVFP4 activations are scaled twice: a static per-tensor global scale frozen + at calibration, and dynamic per-16-block E4M3 scales. Only the block scales + can be derived here, so the global scale is read from a sidecar produced by + an earlier calibration pass (`amax / 6 / 448`, matching the serving + pipeline). Returns ``None`` when no sidecar is configured, leaving block + scales only. + + Keys are resolved through ``act_amax_key_rules`` so modules sharing an + input map to one entry — w1/w3 see the same MoE input and must therefore + share a scale, as must q/k/v. + """ + + path = getattr(self.qcfg, "act_amax_path", None) + if not path: + return None + + table = _load_act_amax_table(path) + + key = self.name + for pattern, replacement in (getattr(self.qcfg, "act_amax_key_rules", None) or []): + if re.search(pattern, self.name): + key = re.sub(pattern, replacement, self.name) + break + + if key not in table: + # Failing loud beats silently falling back to block-only scales, which + # would quantize against a different grid than the kernel uses. + raise KeyError( + f"GPTQ: activation amax for module '{self.name}' (key '{key}') not found in " + f"'{path}'. Add the entry or extend `act_amax_key_rules`." + ) + + amax = float(table[key]) + if not math.isfinite(amax) or amax <= 0.0: + raise ValueError( + f"GPTQ: activation amax for '{key}' must be finite and positive, got {amax}." + ) + + # NVFP4 second-level scale: amax / F4_E2M1_MAX / F8_E4M3_MAX. + return torch.tensor(amax / 6.0 / 448.0, dtype=torch.float32) + def _create_act_quantizer(self) -> Optional[Callable[[torch.Tensor], torch.Tensor]]: """Build the activation fake-quantizer for w4a4, or ``None`` for weight-only. @@ -287,9 +356,18 @@ def _create_act_quantizer(self) -> Optional[Callable[[torch.Tensor], torch.Tenso raise ValueError(f"GPTQ: unsupported `act_format={act_format}`.") block_size = int(getattr(self.qcfg, "act_group_size", 16)) + # Resolved on first use, not here: `self.name` is not final at construction + # (the looper names modules afterwards), and a name-keyed sidecar lookup done + # too early silently reads the wrong entry. + resolved: Dict[str, Optional[torch.Tensor]] = {} def _quantize_act(x: torch.Tensor) -> torch.Tensor: - return fake_quantize_nvfp4(x, block_size=block_size) + if "scale" not in resolved: + resolved["scale"] = self._resolve_act_global_scale() + scale = resolved["scale"] + if scale is not None: + scale = scale.to(x.device) + return fake_quantize_nvfp4(x, block_size=block_size, per_tensor_scale=scale) return _quantize_act diff --git a/tests/test_w4a4_hessian.py b/tests/test_w4a4_hessian.py index 37e50e65c..c9367f0f5 100644 --- a/tests/test_w4a4_hessian.py +++ b/tests/test_w4a4_hessian.py @@ -14,6 +14,7 @@ See `docs/w4a4_hessian.md`. """ +import json import os @@ -130,8 +131,122 @@ def test_hessian_is_deterministic(calib_input): [ (dict(bits=4, group_size=128, format="gptq", act_format=FORMAT.NVFP4), "requires `format=nvfp4`"), (dict(bits=4, group_size=BLOCK, format="nvfp4", act_format=FORMAT.NVFP4, act_group_size=0), "act_group_size"), + (dict(bits=4, group_size=BLOCK, format="nvfp4", act_amax_path="x.json"), "require `act_format`"), ], ) def test_config_rejects_unsupported_activation_settings(kwargs, match): with pytest.raises(ValueError, match=match): QuantizeConfig(**kwargs) + + +# --------------------------------------------------------------- amax sidecar -- + + +def _sidecar(tmp_path, payload) -> str: + path = tmp_path / "amax_per_layer.json" + path.write_text(json.dumps(payload)) + return str(path) + + +def _gptq_named(name: str, qcfg: QuantizeConfig) -> GPTQ: + torch.manual_seed(0) + module = nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False) + gptq = GPTQ(module=module, qcfg=qcfg) + gptq.name = name + return gptq + + +def _qcfg_sidecar(path, rules=None) -> QuantizeConfig: + return QuantizeConfig( + bits=4, group_size=BLOCK, format="nvfp4", sym=True, desc_act=False, + act_format=FORMAT.NVFP4, act_group_size=BLOCK, + act_amax_path=path, act_amax_key_rules=rules, + ) + + +def test_global_scale_matches_pipeline_formula(tmp_path): + """global_scale must be amax / 6 / 448 — the NVFP4 second-level scale.""" + + amax = 12.0 + path = _sidecar(tmp_path, {"layer0.moe_input": amax}) + gptq = _gptq_named("layer0.moe_input", _qcfg_sidecar(path)) + + scale = gptq._resolve_act_global_scale() + assert scale is not None + assert scale.item() == pytest.approx(amax / 6.0 / 448.0) + + +def test_sidecar_accepts_nested_amax_mapping(tmp_path): + """The pipeline emits {'amax': {...}}; a flat mapping must work too.""" + + path = _sidecar(tmp_path, {"amax": {"layer0.moe_input": 6.0}, "meta": {"ignored": 1}}) + gptq = _gptq_named("layer0.moe_input", _qcfg_sidecar(path)) + assert gptq._resolve_act_global_scale().item() == pytest.approx(6.0 / 6.0 / 448.0) + + +def test_key_rules_share_one_scale_across_sibling_modules(tmp_path): + """w1 and w3 consume the same MoE input, so they must resolve to one entry. + + Getting this wrong gives siblings different activation scales than the kernel + applies, which is exactly the mismatch the sidecar exists to prevent. + """ + + path = _sidecar(tmp_path, {"layer7.moe_input": 8.0, "layer7.w2_input": 2.0}) + rules = [ + (r"^layers\.(\d+)\..*\.(?:w1|w3)$", r"layer\1.moe_input"), + (r"^layers\.(\d+)\..*\.w2$", r"layer\1.w2_input"), + ] + cfg = _qcfg_sidecar(path, rules) + + w1 = _gptq_named("layers.7.ffn.experts.3.w1", cfg)._resolve_act_global_scale() + w3 = _gptq_named("layers.7.ffn.experts.3.w3", cfg)._resolve_act_global_scale() + w2 = _gptq_named("layers.7.ffn.experts.3.w2", cfg)._resolve_act_global_scale() + + assert w1.item() == w3.item(), "w1 and w3 share an input and must share a scale" + assert w2.item() != w1.item(), "w2 sees the post-SwiGLU activation and needs its own" + assert w2.item() == pytest.approx(2.0 / 6.0 / 448.0) + + +def test_missing_amax_key_fails_loud(tmp_path): + """A missing entry must raise, not silently fall back to block-only scales.""" + + path = _sidecar(tmp_path, {"layer0.moe_input": 6.0}) + gptq = _gptq_named("layers.99.ffn.experts.0.w1", _qcfg_sidecar(path)) + with pytest.raises(KeyError, match="not found"): + gptq._resolve_act_global_scale() + + +@pytest.mark.parametrize("bad", [0.0, -1.0, float("inf"), float("nan")]) +def test_invalid_amax_rejected(tmp_path, bad): + path = _sidecar(tmp_path, {"k": bad}) + gptq = _gptq_named("k", _qcfg_sidecar(path)) + with pytest.raises(ValueError, match="finite and positive"): + gptq._resolve_act_global_scale() + + +def test_no_sidecar_leaves_global_scale_unset(): + gptq = _gptq_named("anything", _qcfg(act=True)) + assert gptq._resolve_act_global_scale() is None + + +def test_global_scale_changes_the_hessian(tmp_path, calib_input): + """The sidecar must actually reach the quantizer, not just parse.""" + + torch.manual_seed(0) + inp = calib_input + + def hessian(cfg, name="k"): + torch.manual_seed(0) + g = GPTQ(module=nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False), qcfg=cfg) + g.name = name + g._act_quantizer = g._create_act_quantizer() + g.add_batch(inp, torch.empty(0)) + return g.finalize_hessian().clone() + + block_only = hessian(_qcfg(act=True)) + # An amax far from the data's own range forces a visibly different grid. + with_scale = hessian(_qcfg_sidecar(_sidecar(tmp_path, {"k": 40.0}))) + + assert not torch.allclose(block_only, with_scale), ( + "global scale had no effect on the Hessian — sidecar is not reaching the quantizer" + ) From 64667668f4b17679ecabc8773080445549490ab8 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 18:06:56 -0700 Subject: [PATCH 05/14] fix(nvfp4): add the two-level weight_scale_2 and match torchao's numerics GPTQModel's NVFP4 export stored only per-block E4M3 scales. ModelOpt, Quark and tore-quant all use two-level scaling: a per-tensor weight_scale_2 = amax/(6*448) plus block scales. Without it the block scales carry absolute magnitudes, and for realistic weight tensors block_amax/6 sits below E4M3's smallest normal (2^-6) -- measured 100% of blocks clamped at the floor on an N(0, 0.02) weight -- so every block quantized on a compressed grid. Two-level round-trip error is 1.36x lower on that tensor (0.096 vs 0.131 relative). RTN path: quantize_nvfp4_weight computes the canonical greedy global scale (verified equal to ModelOpt's global_amax/(E2M1_MAX*m_fp8) and tore-quant's amax_to_scale) and returns a 3-tuple. GPTQ path: the loop has already fixed each block scale s_i and the packed codes are encoded against them, so the split must satisfy block_e4m3_i * ws2 == s_i exactly. A greedy ws2 cannot (the ratio of two E4M3 mantissas is generally unrepresentable); a power-of-two ws2 = 2^ceil(log2(max(s)/448)) only shifts exponents, leaving mantissas intact. Verified bit-exact recovery; an equality check raises rather than silently emitting codes that no longer match their scales. Both read paths apply it: dequantize_weight folds ws2 into the block scales (one pass, exact in fp32), _native_weight passes it as NVFP4Tensor's per_tensor_scale. TorchFP4Linear -- which executes directly from checkpoint weights -- now accepts weight_scale_2 too; it previously ignored the tensor that ModelOpt-layout checkpoints (e.g. MiniMax-M3-NVFP4) carry on every quantized module, dequantizing every block off by a constant factor. Also switch x/s to x * s.reciprocal() in NVFP4Quantizer.quantize and _pack_with_gptq_scales, matching torchao's nvfp4_quantize ("multiply by reciprocal instead of dividing to match MSLK triton kernel numerics"). The 1-ULP difference can flip an E2M1 code on rounding boundaries, which breaks bit-identical packing claims. An unpacked module defaults weight_scale_2 to 1.0 (identity), so one-level checkpoints load unchanged. test_direct_pack_matches_torchao updated to assert the two-level layout. Co-Authored-By: Claude Fable 5 --- gptqmodel/nn_modules/qlinear/fp4.py | 110 +++++++++++++++++++++++++--- gptqmodel/quantization/quantizer.py | 8 +- tests/test_nvfp4_export.py | 13 +++- 3 files changed, 118 insertions(+), 13 deletions(-) diff --git a/gptqmodel/nn_modules/qlinear/fp4.py b/gptqmodel/nn_modules/qlinear/fp4.py index 25875e88f..1a3ab5756 100644 --- a/gptqmodel/nn_modules/qlinear/fp4.py +++ b/gptqmodel/nn_modules/qlinear/fp4.py @@ -43,6 +43,7 @@ def __init__( weight_block_size: int, orig_dtype: torch.dtype, bias: Optional[torch.Tensor] = None, + weight_scale_2: Optional[torch.Tensor] = None, ) -> None: super().__init__() self.in_features = in_features @@ -51,6 +52,12 @@ def __init__( self.orig_dtype = orig_dtype self.register_buffer("weight", weight) self.register_buffer("weight_scale", weight_scale) + # ModelOpt/Quark NVFP4 checkpoints carry a second-level per-tensor scale. + # Ignoring it would dequantize every block by a constant factor off. + if isinstance(weight_scale_2, torch.Tensor): + self.register_buffer("weight_scale_2", weight_scale_2.to(torch.float32).reshape(())) + else: + self.weight_scale_2 = None if isinstance(bias, torch.Tensor): self.register_buffer("bias", bias) else: @@ -66,6 +73,7 @@ def _native_weight(self) -> "NVFP4Tensor": self.weight_scale, block_size=self.weight_block_size, orig_dtype=self.orig_dtype, + per_tensor_scale=self.weight_scale_2, ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -77,12 +85,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return F.linear(x, self._native_weight(), bias) +# NVFP4's second-level (per-tensor) scale. ModelOpt, Quark and tore-quant all +# derive it the same way -- amax / (E2M1_MAX * F8E4M3_MAX) -- and serialize it as +# `weight_scale_2`. Without it the E4M3 block scales carry absolute magnitudes: +# for typical weights `block_amax/6` lands below E4M3's smallest normal (2^-6), +# so `clamp(min=E4M3_EPS)` floors them and the block quantizes on a compressed +# grid. With it, block scales become 448*block_amax/global_amax and use E4M3's +# full range. +_F4_E2M1_MAX = 6.0 +_F8_E4M3_MAX = 448.0 + + +def compute_nvfp4_global_scale(weight: torch.Tensor) -> torch.Tensor: + """NVFP4 `weight_scale_2` for one tensor: amax / (6 * 448).""" + + amax = weight.detach().abs().max().to(torch.float32) + scale = amax / (_F4_E2M1_MAX * _F8_E4M3_MAX) + # An all-zero tensor would give 0 and make every block scale non-finite. + return torch.where(scale > 0, scale, torch.ones_like(scale)).reshape(()) + + def quantize_nvfp4_weight( weight: torch.Tensor, *, weight_block_size: int = NVFP4_WEIGHT_BLOCK_SIZE, -) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize one dense 2D weight to NVFP4 packed bytes plus per-block E4M3 scales.""" +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize one dense 2D weight to NVFP4: packed bytes, E4M3 block scales, global scale.""" if nvfp4_quantize is None: raise RuntimeError("NVFP4 quantization requires torchao with NVFP4 support.") @@ -97,12 +125,15 @@ def quantize_nvfp4_weight( ) weight = weight.to(device="cpu", dtype=torch.float32).contiguous() - scales, packed = nvfp4_quantize(weight, block_size=NVFP4_WEIGHT_BLOCK_SIZE) + global_scale = compute_nvfp4_global_scale(weight) + scales, packed = nvfp4_quantize( + weight, block_size=NVFP4_WEIGHT_BLOCK_SIZE, per_tensor_scale=global_scale + ) # Store packed nibbles as uint8 so safetensors serialization matches the # ModelOpt/torchao NVFP4 checkpoint layout regardless of torch fp4 dtype support. if packed.dtype is not torch.uint8: packed = packed.view(torch.uint8) - return packed.contiguous(), scales.contiguous() + return packed.contiguous(), scales.contiguous(), global_scale.contiguous() class TorchNVFP4Linear(WeightOnlyQuantLinear): @@ -191,6 +222,14 @@ def _allocate_buffers(self, *, bias: bool) -> None: else: self.register_buffer("weight_scale", scale) + # NVFP4's second-level per-tensor scale. 1.0 is the identity, so an + # un-packed module still dequantizes correctly. + scale_2 = torch.ones((), dtype=torch.float32) + if "weight_scale_2" in self._buffers: + self.weight_scale_2 = scale_2 + else: + self.register_buffer("weight_scale_2", scale_2) + if bias: bias_tensor = torch.zeros(self.out_features, dtype=torch.float16) if "bias" in self._buffers: @@ -206,6 +245,8 @@ def list_buffers(self): buffers.append(self.weight) if hasattr(self, "weight_scale") and self.weight_scale is not None: buffers.append(self.weight_scale) + if hasattr(self, "weight_scale_2") and self.weight_scale_2 is not None: + buffers.append(self.weight_scale_2) if hasattr(self, "bias") and self.bias is not None: buffers.append(self.bias) return buffers @@ -247,12 +288,21 @@ def pack_gpu( del block_in, device self.pack_original(linear=linear, scales=scales, zeros=zeros, g_idx=g_idx) - def _pack_with_gptq_scales(self, weight: torch.Tensor, scales: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def _pack_with_gptq_scales( + self, weight: torch.Tensor, scales: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Pack a weight already on the NVFP4 grid using the loop-produced group scales. GPTQ error feedback mutates columns after each group scale is fixed, so the final weight must be encoded against those exact scales — re-deriving scales from the weight (RTN-style) would land some blocks on a different grid. + + That constraint also dictates how `weight_scale_2` is chosen here. The loop + already fixed each block scale `s_i`, so the two-level form must satisfy + `block_e4m3_i * ws2 == s_i` exactly, or the packed codes stop matching their + scales. A **power-of-two** `ws2` makes that exact: dividing by it only shifts + the exponent, leaving the E4M3 mantissa untouched. A greedy `amax/(6*448)` + would not, so it is used only on the RTN path where scales are free. """ if f32_to_f4_unpacked is None or pack_uint4 is None: @@ -267,12 +317,35 @@ def _pack_with_gptq_scales(self, weight: torch.Tensor, scales: torch.Tensor) -> ) expanded = scale.repeat_interleave(self.weight_block_size, dim=1) - scaled = torch.clamp(weight / expanded, -6.0, 6.0) + # Reciprocal-multiply, not divide: matches torchao's nvfp4_quantize numerics + # (see NVFP4Quantizer.quantize). A 1-ULP difference can flip an E2M1 code. + scaled = torch.clamp(weight * expanded.reciprocal(), -6.0, 6.0) codes = f32_to_f4_unpacked(scaled.contiguous()) packed = pack_uint4(codes) if packed.dtype is not torch.uint8: packed = packed.view(torch.uint8) - return packed.contiguous(), scale.to(torch.float8_e4m3fn).contiguous() + + # ws2 = 2^ceil(log2(max(s) / 448)): the smallest power of two that brings + # every block scale within E4M3's range. + smax = scale.max() + if smax > 0: + exponent = torch.ceil(torch.log2(smax / _F8_E4M3_MAX)) + global_scale = torch.exp2(exponent).to(torch.float32).reshape(()) + else: + global_scale = torch.ones((), dtype=torch.float32) + + block_scale = (scale * global_scale.reciprocal()).to(torch.float8_e4m3fn) + # The exactness the docstring depends on is cheap to verify, and a silent + # mismatch here means every packed code is read against the wrong scale. + recovered = block_scale.to(torch.float32) * global_scale + if not torch.equal(recovered, scale): + bad = int((recovered != scale).sum()) + raise RuntimeError( + f"NVFP4 two-level split changed {bad} of {scale.numel()} block scales; " + "packed codes would no longer match their scales." + ) + + return packed.contiguous(), block_scale.contiguous(), global_scale.contiguous() @torch.inference_mode() def pack_original( @@ -288,14 +361,14 @@ def pack_original( weight = self._weight_to_matrix(linear).to(device="cpu", dtype=torch.float32) if isinstance(scales, torch.Tensor) and scales.numel() > 0: - qweight, weight_scale = self._pack_with_gptq_scales(weight, scales) + qweight, weight_scale, weight_scale_2 = self._pack_with_gptq_scales(weight, scales) else: weight = _apply_optional_smoother( weight, smooth=smooth, group_size=self.smooth_block_size(), ) - qweight, weight_scale = quantize_nvfp4_weight( + qweight, weight_scale, weight_scale_2 = quantize_nvfp4_weight( weight, weight_block_size=self.weight_block_size, ) @@ -310,6 +383,11 @@ def pack_original( else: self.register_buffer("weight_scale", weight_scale) + if "weight_scale_2" in self._buffers: + self.weight_scale_2 = weight_scale_2 + else: + self.register_buffer("weight_scale_2", weight_scale_2) + if linear.bias is not None: bias = linear.bias.detach().to(device="cpu", dtype=torch.float16) if "bias" in self._buffers: @@ -339,6 +417,12 @@ def dequantize_weight( if self.weight_scale.device == target_device else self.weight_scale.to(device=target_device) ) + # Fold the per-tensor scale into the block scales rather than scaling the + # dequantized weight afterwards: one pass, and it keeps the fp32 product + # exact instead of rounding through the E4M3 block scale twice. + scale_2 = getattr(self, "weight_scale_2", None) + if isinstance(scale_2, torch.Tensor): + scale = scale.to(dtype=torch.float32) * scale_2.to(device=target_device, dtype=torch.float32) weight = dequantize_f4_e2m1( source, scale=scale, @@ -357,6 +441,7 @@ def _native_weight(self) -> "NVFP4Tensor": self.weight_scale, block_size=self.weight_block_size, orig_dtype=torch.bfloat16, + per_tensor_scale=getattr(self, "weight_scale_2", None), ) def _can_use_native_linear(self, x_flat: torch.Tensor) -> bool: @@ -397,4 +482,9 @@ def forward(self, x: torch.Tensor): return output.reshape(original_shape) -__all__ = ["TorchFP4Linear", "TorchNVFP4Linear", "quantize_nvfp4_weight"] +__all__ = [ + "TorchFP4Linear", + "TorchNVFP4Linear", + "quantize_nvfp4_weight", + "compute_nvfp4_global_scale", +] diff --git a/gptqmodel/quantization/quantizer.py b/gptqmodel/quantization/quantizer.py index f266857d8..3ac5f3a55 100644 --- a/gptqmodel/quantization/quantizer.py +++ b/gptqmodel/quantization/quantizer.py @@ -215,7 +215,13 @@ def find_params(self, x, weight=False): def quantize(self, x): orig_dtype = x.dtype - scaled = torch.clamp(x.to(torch.float32) / self.scale, -_F4_E2M1_MAX, _F4_E2M1_MAX) + # Multiply by the reciprocal rather than dividing, matching torchao's + # nvfp4_quantize (and the MSLK triton kernel it tracks). `x / s` and + # `x * (1/s)` differ by up to 1 ULP in fp32, and with only 16 E2M1 codes a + # value on a rounding boundary can land on a different code. + scaled = torch.clamp( + x.to(torch.float32) * self.scale.reciprocal(), -_F4_E2M1_MAX, _F4_E2M1_MAX + ) codes = f32_to_f4_unpacked(scaled.contiguous()) dequant = f4_unpacked_to_f32(codes) * self.scale return dequant.to(orig_dtype) diff --git a/tests/test_nvfp4_export.py b/tests/test_nvfp4_export.py index 72c9ba820..6d534ecd1 100644 --- a/tests/test_nvfp4_export.py +++ b/tests/test_nvfp4_export.py @@ -157,15 +157,23 @@ def test_direct_pack_matches_torchao(): out_features=64, bias=True, register_buffers=True, + """RTN pack must match torchao's TWO-level path: per-tensor scale amax/(6*448) + plus E4M3 block scales, matching the ModelOpt/Quark/tore-quant layout.""" + ) module.pack_original(linear=lin, scales=None, zeros=None) - scales_ref, packed_ref = nvfp4_quantize(lin.weight.detach().float(), block_size=BLOCK) + w = lin.weight.detach().float() + gs_ref = (w.abs().max() / (6.0 * 448.0)).reshape(()) + scales_ref, packed_ref = nvfp4_quantize(w.contiguous(), block_size=BLOCK, per_tensor_scale=gs_ref) assert torch.equal(module.weight, packed_ref.view(torch.uint8)) assert torch.equal(module.weight_scale.view(torch.uint8), scales_ref.view(torch.uint8)) deq = module.dequantize_weight(dtype=torch.float32).T - assert torch.allclose(deq, _nvfp4_dequant(packed_ref, scales_ref), atol=1e-3) + ref = NVFP4Tensor( + packed_ref, scales_ref, BLOCK, torch.float32, per_tensor_scale=gs_ref + ).dequantize(torch.float32) + assert torch.allclose(deq, ref, atol=1e-3) def test_gptq_scale_pack_is_lossless(): @@ -175,6 +183,7 @@ def test_gptq_scale_pack_is_lossless(): torch.manual_seed(2) IN, OUT = 128, 64 qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + assert torch.equal(module.weight_scale_2, gs_ref) quantizer = NVFP4Quantizer(qcfg=qcfg) quantizer.configure(perchannel=True) From 3c476d280e0880485e02b42eff751d53aee4d6af Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 18:12:35 -0700 Subject: [PATCH 06/14] feat(scripts): minimax-m3 NVFP4 w4a4 pipeline -- amax sidecar, quantize driver, run.sh Three-stage pipeline for quantizing Minimax-M3-0602 routed experts to NVFP4 with w4a4 calibration: emit_act_amax.py (step 0) -- emits the per-layer activation amax sidecar the static NVFP4 global scale needs. Two keys per MoE layer, keyed by input site because modules sharing an input share a scale: layer{N}.moe_input (gate/up_proj) and layer{N}.w2_input (down_proj, post-SwiGLU). Plain max over |x| reduced across a layer's experts, matching tore-quant's save_amax_sidecar peer-max semantics. Captures via down_proj pre-hooks when the loader exposes per-expert Linears (GPTQModel does: it splits the checkpoint's fused 3D expert tensors), falling back to wrapping the fused module's _apply_gate -- whose return value IS the down_proj input -- and refusing to emit a partial sidecar if neither works. Streams layer windows through one GPU with per-window resume, the p1b_expert_coverage harness pattern. quantize_w4a4.py -- GPTQ->NVFP4 driver. Quantizes only routed experts (57 layers x 128 experts x 3 projections); attention, shared experts, dense MLP 0-2, routers, embeddings and the vision tower pass through, matching the tore-quant recipe format map. RTN-fallback thresholds are per-projection at each weight's in_features (gate/up 6144, down 3072): H = XtX has rank <= min(n_tokens, cols), so a module below its own input dim would be GPTQ-solved on a singular Hessian propped up by damping. The library's 0.5% default sits below both dims and is scale-invariant besides (RTN iff routing share < pct/(100*top_k) -- more calibration data cannot lower it). With ACT_AMAX set, the Hessian is accumulated on Q(X) with the sidecar-derived static global scale; without it the run is weight-only and says so. Post-check asserts the saved layout (uint8 weight + e4m3 weight_scale + fp32 weight_scale_2) instead of trusting the log. run.sh -- preflight (reads file CONTENT, not stat: the /scratch checkpoint copy passes ls and fails on read), disk/GPU checks, predicted per-projection RTN rates from the measured coverage matrix, then the driver. All knobs are env vars; PREFLIGHT_ONLY=1 and LAYERS=N smoke-test modes included. Co-Authored-By: Claude Fable 5 --- scripts/minimax_m3_nvfp4/emit_act_amax.py | 276 ++++++++++++++++++++++ scripts/minimax_m3_nvfp4/quantize_w4a4.py | 200 ++++++++++++++++ scripts/minimax_m3_nvfp4/run.sh | 143 +++++++++++ 3 files changed, 619 insertions(+) create mode 100644 scripts/minimax_m3_nvfp4/emit_act_amax.py create mode 100644 scripts/minimax_m3_nvfp4/quantize_w4a4.py create mode 100755 scripts/minimax_m3_nvfp4/run.sh diff --git a/scripts/minimax_m3_nvfp4/emit_act_amax.py b/scripts/minimax_m3_nvfp4/emit_act_amax.py new file mode 100644 index 000000000..9e108bf07 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/emit_act_amax.py @@ -0,0 +1,276 @@ +"""Step 0 for w4a4: emit the per-layer activation amax sidecar. + +NVFP4 activations carry a *static* per-tensor global scale (calibrated, frozen at +serving) plus dynamic per-16-block E4M3 scales. Only the block scales can be +derived while GPTQ accumulates its Hessian, so the global scale has to come from +an earlier pass -- this one. Consumed via `ACT_AMAX=` in run.sh, which turns it +into `global_scale = amax / 6 / 448` (see docs/w4a4_hessian.md). + +Scales are keyed by *input site*, not by module, because modules sharing an input +share a scale: + + layer{N}.moe_input -> input to routed-expert gate_proj/up_proj (w1/w3) + layer{N}.w2_input -> input to routed-expert down_proj (w2), post-SwiGLU + +A plain max is taken over abs, matching tore-quant's `save_amax_sidecar` +(no percentile or MSE clipping), and reduced across every expert of a layer -- +the peer-max sync `fixup_moe_expert_amax` performs. + +Streams the model window-by-window through one GPU exactly as +p1b_expert_coverage.py does (materialize -> forward all samples -> free), since +the checkpoint is far larger than any single device. +""" + +import json +import os +import random +import time + +import torch +import torch.nn as nn +from transformers import AutoTokenizer + +MODEL = os.environ.get("MODEL", "/data/huggingface/Minimax-M3-0602") +OUT_JSON = os.environ.get( + "OUT_JSON", "/data/sgambhira/GPTQModel/scripts/minimax_m3_nvfp4/amax_per_layer.json" +) +SNAP = os.environ.get( + "SNAP", + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745", +) +OFFLOAD = os.environ.get("OFFLOAD", "/data/sgambhira/gptqmodel-offload/amax") +SEED = int(os.environ.get("SEED", "42")) +NUM_SAMPLES = int(os.environ.get("NUM_SAMPLES", "64")) +MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "8192")) +WINDOW = int(os.environ.get("WINDOW", "5")) +DEVICE = torch.device(os.environ.get("DEVICE", "cuda:0")) +CKPT = os.environ.get("CKPT", "/data/sgambhira/gptqmodel-offload/amax_resume.pt") + +t0 = time.time() + + +def log(msg: str) -> None: + print(f"[amax +{time.time() - t0:7.1f}s] {msg}", flush=True) + + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4", offload_to_disk_path=OFFLOAD) +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +log("model loaded") + +lm = model.model.model.language_model +all_layers = list(lm.layers) +num_layers = len(all_layers) + +# --- calibration inputs (identical to the coverage / GPTQ calibration set) --- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(line) for line in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) +inputs = [] +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + ids = tokenizer(text, add_special_tokens=False, return_tensors="pt")["input_ids"][:, :MAX_TOKENS] + inputs.append(ids) +total_tokens = sum(i.shape[1] for i in inputs) +log(f"calibration: {len(inputs)} samples, {total_tokens:,} tokens") + +# --- amax capture ----------------------------------------------------------- +amax: dict[str, float] = {} + + +def _update(key: str, tensor: torch.Tensor) -> None: + if not torch.is_tensor(tensor) or tensor.numel() == 0: + return + v = tensor.detach().abs().max().to(torch.float32).item() + if v > amax.get(key, 0.0): + amax[key] = v + + +def make_moe_input_hook(layer_idx: int): + # Pre-hook on the experts module: args[0] is the hidden state entering w1/w3. + def hook(module, args, kwargs): + if args and torch.is_tensor(args[0]): + _update(f"layer{layer_idx}.moe_input", args[0]) + return hook + + +def make_apply_gate_patch(layer_idx: int, original): + """Capture w2_input by wrapping the fused experts' `_apply_gate`. + + MiniMaxM3VLExperts stores experts as 3D nn.Parameters -- `gate_up_proj` and + `down_proj` -- and runs them with F.linear inside a Python loop, so there is + no down_proj submodule to hook. Its forward is: + + current = self._apply_gate(F.linear(hidden[idx], gate_up_proj[e])) + current = F.linear(current, down_proj[e]) * top_k_weights[...] + + so `_apply_gate`'s return value *is* the down_proj input -- the post-SwiGLU + activation we need. Wrapping it captures that exactly, with no + reimplementation of the SwiGLU to drift out of sync. It fires once per hit + expert, so the max is already reduced across the layer's experts (the + peer-max sync). + """ + + def patched(gate_up: torch.Tensor) -> torch.Tensor: + out = original(gate_up) + _update(f"layer{layer_idx}.w2_input", out) + return out + + return patched + + +def make_w2_input_hook(layer_idx: int): + # Pre-hook on a per-expert down_proj: its input IS the post-SwiGLU activation. + # Fires once per expert, so the max is already reduced across the layer. + def hook(module, args): + if args and torch.is_tensor(args[0]): + _update(f"layer{layer_idx}.w2_input", args[0]) + return hook + + +# The experts tree differs by loader, so support both rather than betting on one: +# * GPTQModel.load splits the checkpoint's fused 3D tensors into 128 per-expert +# nn.Linear modules (gate_proj / up_proj / down_proj) -- hook down_proj. +# * Stock transformers keeps MiniMaxM3VLExperts fused as 3D nn.Parameters with +# no submodules -- wrap `_apply_gate`, whose return value is the down_proj +# input. +hook_handles = [] +patched_experts = [] +moe_layers = [] +n_down = 0 +for li, layer in enumerate(all_layers): + experts = getattr(getattr(layer, "mlp", None), "experts", None) + if not isinstance(experts, nn.Module): + continue + moe_layers.append(li) + hook_handles.append( + experts.register_forward_pre_hook(make_moe_input_hook(li), with_kwargs=True) + ) + + downs = [m for name, m in experts.named_modules() + if name.endswith("down_proj") and isinstance(m, nn.Linear)] + if downs: + for m in downs: + hook_handles.append(m.register_forward_pre_hook(make_w2_input_hook(li))) + n_down += len(downs) + continue + + original_apply_gate = getattr(experts, "_apply_gate", None) + if original_apply_gate is None: + raise SystemExit( + f"layer {li}: {type(experts).__name__} exposes neither down_proj Linears nor " + "`_apply_gate`; w2_input cannot be captured. Adapt the capture to this " + "architecture rather than emitting a sidecar missing every down_proj scale." + ) + # Instance attribute shadows the bound method, so `self._apply_gate(x)` hits it. + experts._apply_gate = make_apply_gate_patch(li, original_apply_gate) + patched_experts.append(experts) + +log(f"hooked {len(moe_layers)} MoE layers ({moe_layers[:3]}...{moe_layers[-2:]}): " + f"moe_input via experts pre-hook; w2_input via {n_down} down_proj hooks " + f"+ {len(patched_experts)} _apply_gate wrappers") + +# --- windowed streaming forward (same structure as p1b_expert_coverage.py) --- +model.shell_module_materialize(lm.rotary_emb, DEVICE) +model.shell_module_materialize(lm.norm, DEVICE) + +start_layer = 0 +if os.path.exists(CKPT): + state = torch.load(CKPT, map_location="cpu", weights_only=False) + hidden = state["hidden"] + amax.update(state["amax"]) + start_layer = state["next_layer"] + log(f"resumed from checkpoint at layer {start_layer}") +else: + embed = model.shell_module_materialize(lm.embed_tokens, DEVICE) + hidden = [] + with torch.inference_mode(): + for ids in inputs: + hidden.append(embed(ids.to(DEVICE)).to("cpu")) + log("embedded all samples") + +orig_layers = lm.layers + +for w_start in range(start_layer, num_layers, WINDOW): + w_layers = all_layers[w_start : w_start + WINDOW] + lm.layers = orig_layers + for layer in w_layers: + model.shell_module_materialize(layer, DEVICE) + lm.layers = nn.ModuleList(w_layers) + + window_out = {} + + def tail_hook(module, args, output): + window_out["h"] = (output[0] if isinstance(output, tuple) else output).detach() + + tail = w_layers[-1].register_forward_hook(tail_hook) + + with torch.inference_mode(): + for si in range(len(hidden)): + h = hidden[si].to(DEVICE) + pos = torch.arange(h.shape[1], device=DEVICE).unsqueeze(0) + lm(inputs_embeds=h, position_ids=pos, use_cache=False) + hidden[si] = window_out["h"].to("cpu") + + tail.remove() + for layer in w_layers: + layer.to_empty(device="meta") + torch.cuda.empty_cache() + log(f"window done: layers {w_start}-{w_start + len(w_layers) - 1}") + + next_layer = w_start + len(w_layers) + if next_layer < num_layers: + tmp = CKPT + ".tmp" + torch.save({"hidden": hidden, "amax": amax, "next_layer": next_layer}, tmp) + os.replace(tmp, CKPT) + +lm.layers = orig_layers +for h in hook_handles: + h.remove() +for experts in patched_experts: + # Drop the instance attribute so the class's bound method is live again. + del experts._apply_gate + +# --- emit ------------------------------------------------------------------- +missing = [ + f"layer{li}.{site}" + for li in moe_layers + for site in ("moe_input", "w2_input") + if f"layer{li}.{site}" not in amax +] +if missing: + raise SystemExit(f"missing amax entries, refusing to write a partial sidecar: {missing[:8]}") + +payload = { + "amax": {k: amax[k] for k in sorted(amax, key=lambda s: (int(s.split(".")[0][5:]), s))}, + "meta": { + "model": MODEL, + "samples": NUM_SAMPLES, + "max_tokens": MAX_TOKENS, + "seed": SEED, + "total_tokens": total_tokens, + "method": "plain max over |x|, reduced across experts (peer-max)", + "sites": {"moe_input": "gate_proj/up_proj input", "w2_input": "down_proj input"}, + }, +} +os.makedirs(os.path.dirname(OUT_JSON) or ".", exist_ok=True) +with open(OUT_JSON, "w") as f: + json.dump(payload, f, indent=1) +log(f"wrote {OUT_JSON} ({len(payload['amax'])} entries)") + +if os.path.exists(CKPT): + os.remove(CKPT) + +vals = list(payload["amax"].values()) +log(f"amax range: min={min(vals):.4g} max={max(vals):.4g}") +log(f"implied global_scale range: {min(vals) / 6 / 448:.4g} .. {max(vals) / 6 / 448:.4g}") +for li in moe_layers[:3] + moe_layers[-2:]: + log(f" layer{li}: moe_input={amax[f'layer{li}.moe_input']:.4g} " + f"w2_input={amax[f'layer{li}.w2_input']:.4g}") +log("DONE") diff --git a/scripts/minimax_m3_nvfp4/quantize_w4a4.py b/scripts/minimax_m3_nvfp4/quantize_w4a4.py new file mode 100644 index 000000000..fe98ec88d --- /dev/null +++ b/scripts/minimax_m3_nvfp4/quantize_w4a4.py @@ -0,0 +1,200 @@ +"""GPTQ -> NVFP4 quantization of Minimax-M3-0602, calibrated for w4a4. + +Quantizes the routed experts of all 57 MoE layers to NVFP4. Everything else -- +attention, index projections, shared experts, the 3 dense MLP layers, routers, +norms, embeddings, lm_head, vision tower -- is left untouched and passes through +in the source checkpoint's precision, matching the format map in +tore-quant/recipes/minimax-m3-nvfp4/README.md. + +When an activation amax sidecar is supplied, the GPTQ Hessian is accumulated on +Q(X) rather than X, so the weights are solved for the activations a w4a4 kernel +actually consumes (docs/w4a4_hessian.md). Without one it degrades to a +weight-only calibration and says so. + +Driven by run.sh; all knobs are environment variables. +""" + +import json +import os +import random +import sys +import time + +import torch +from transformers import AutoTokenizer + +MODEL = os.environ.get("MODEL", "/data/huggingface/Minimax-M3-0602") +OUT = os.environ.get("OUT", "/data/huggingface/Minimax-M3-0602-NVFP4-GPTQ-w4a4") +SNAP = os.environ.get( + "SNAP", + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745", +) +OFFLOAD = os.environ.get("OFFLOAD", "/data/sgambhira/gptqmodel-offload/w4a4") +ACT_AMAX = os.environ.get("ACT_AMAX", "").strip() +SEED = int(os.environ.get("SEED", "42")) +NUM_SAMPLES = int(os.environ.get("NUM_SAMPLES", "64")) +MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "8192")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "1")) +# Fallback threshold, in TOKENS seen by a module (not samples). +# +# The library default "0.5%" is a fixed quantile of the routing distribution and +# is scale-invariant -- an expert is RTN'd iff its routing share < pct/(100*top_k), +# with total tokens cancelling -- so more calibration data does not reduce it. +# Worse, at 2,355 tokens it sits below the input dim, so ~22% of the modules it +# lets through are GPTQ-solved on a singular Hessian held up only by damping. +# +# H = XᵀX is (columns x columns) and rank(XᵀX) = rank(X) <= min(n_tokens, columns), +# so a module needs at least `columns` (= in_features) tokens to be full rank. +# Default here is per-projection thresholds at exactly that bound. +FALLBACK_THRESHOLD = os.environ.get("FALLBACK_THRESHOLD", "").strip() +# Multiplier on in_features. 1.0 = bare full rank (well-defined but ill-conditioned +# at exactly n == columns); >1 buys conditioning margin at the cost of more RTN. +RANK_MARGIN = float(os.environ.get("RANK_MARGIN", "1.0")) +LAYERS = os.environ.get("LAYERS", "").strip() # e.g. "3" or "3-5"; empty = all + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 +from gptqmodel.quantization.config import FORMAT, Fallback, FallbackStrategy # noqa: E402 + +t0 = time.time() + + +def log(msg: str) -> None: + print(f"[w4a4 +{time.time() - t0:7.1f}s] {msg}", flush=True) + + +# --- what to quantize ------------------------------------------------------- +# Only routed experts. `-:` prefixes a skip rule. +dynamic = { + r"-:.*self_attn\..*": {}, + r"-:.*shared_experts\..*": {}, + r"-:.*\.layers\.(0|1|2)\.mlp\..*": {}, + r"-:.*mlp\.gate\..*": {}, # routers stay F32 + r"-:.*(embed_tokens|lm_head).*": {}, + r"-:.*(vision|visual|mm_projector|patch_merge).*": {}, +} +if LAYERS: + lo, _, hi = LAYERS.partition("-") + keep = set(range(int(lo), int(hi or lo) + 1)) + skip = [str(i) for i in range(60) if i not in keep] + dynamic[rf"-:.*\.layers\.({'|'.join(skip)})\..*"] = {} + log(f"LAYERS={LAYERS}: quantizing only {sorted(keep)}") + +# --- RTN fallback threshold, keyed to each projection's Hessian rank --------- +# All three projections of one expert see the same routed tokens, but they have +# different in_features, so a single global threshold either admits singular +# solves (too low) or RTNs down_proj modules that were perfectly well-posed +# (too high). `dynamic` supports a per-module `fallback` override, so key it to +# the actual dimension: gate/up_proj are (3072, 6144) and down_proj (6144, 3072). +IN_FEATURES = {"gate_proj": 6144, "up_proj": 6144, "down_proj": 3072} + +if FALLBACK_THRESHOLD: + default_fallback = Fallback(strategy=FallbackStrategy.RTN, threshold=FALLBACK_THRESHOLD) + log(f"fallback: global threshold {FALLBACK_THRESHOLD} (overriding rank-based default)") +else: + # Anything not matched below keeps the strictest bound. + default_fallback = Fallback(strategy=FallbackStrategy.RTN, + threshold=int(max(IN_FEATURES.values()) * RANK_MARGIN)) + for proj, cols in IN_FEATURES.items(): + thr = int(cols * RANK_MARGIN) + dynamic[rf".*mlp\.experts\.\d+\.{proj}$"] = { + "fallback": {"strategy": "rtn", "threshold": thr} + } + log("fallback: per-projection rank thresholds " + + ", ".join(f"{p}={int(c * RANK_MARGIN)}" for p, c in IN_FEATURES.items()) + + f" (RANK_MARGIN={RANK_MARGIN}) -- every GPTQ solve gets a full-rank Hessian") + +# --- activation quantization (w4a4) ---------------------------------------- +act_kwargs = {} +if ACT_AMAX: + if not os.path.isfile(ACT_AMAX): + sys.exit(f"ACT_AMAX sidecar not found: {ACT_AMAX}") + act_kwargs = dict( + act_format=FORMAT.NVFP4, + act_group_size=16, + act_amax_path=ACT_AMAX, + # w1/w3 (gate_proj/up_proj) share the MoE input, so they share one amax + # entry; w2 (down_proj) sees the post-SwiGLU activation and gets its own. + act_amax_key_rules=[ + (r".*\.layers\.(\d+)\.mlp\.experts\.\d+\.(?:gate_proj|up_proj)$", r"layer\1.moe_input"), + (r".*\.layers\.(\d+)\.mlp\.experts\.\d+\.down_proj$", r"layer\1.w2_input"), + ], + ) + log(f"w4a4: Hessian will be accumulated on Q(X), amax sidecar {ACT_AMAX}") +else: + log("NO ACT_AMAX: weight-only calibration (Hessian on clean X). NOT w4a4-optimal.") + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + sym=True, + desc_act=False, + dynamic=dynamic, + fallback=default_fallback, + moe_vram_strategy="balanced", + offload_to_disk_path=OFFLOAD, + **act_kwargs, +) +log(f"qcfg format={qcfg.format} bits={qcfg.bits} group_size={qcfg.group_size}") + +# --- calibration: 64 seed-42 HLE trajectories, first 8k tokens -------------- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(line) for line in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) + +calibration = [] +total_tokens = 0 +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + enc = tokenizer(text, add_special_tokens=False, return_tensors="pt") + calibration.append({ + "input_ids": enc["input_ids"][:, :MAX_TOKENS], + "attention_mask": enc["attention_mask"][:, :MAX_TOKENS], + }) + total_tokens += calibration[-1]["input_ids"].shape[1] + +log(f"calibration: {len(calibration)} samples, {total_tokens:,} tokens") + +# --- run -------------------------------------------------------------------- +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +log("model loaded") + +model.quantize(calibration, batch_size=BATCH_SIZE, backend=BACKEND.TORCH) +log("quantize complete") + +model.save(OUT) +log(f"saved -> {OUT}") + +# --- post-check: the export must actually be NVFP4 -------------------------- +# A silently-passed-through expert is the failure mode that looks like success, +# so assert the layout rather than trusting the run log. +from safetensors import safe_open # noqa: E402 + +index = json.load(open(f"{OUT}/model.safetensors.index.json"))["weight_map"] +probe = next((k for k in index if ".mlp.experts.0.down_proj.weight" in k), None) +if probe is None: + sys.exit("post-check FAILED: no routed-expert weight found in the saved index") + +with safe_open(f"{OUT}/{index[probe]}", framework="pt", device="cpu") as f: + w = f.get_tensor(probe) +scale_key = probe.replace(".weight", ".weight_scale") +if scale_key not in index: + sys.exit(f"post-check FAILED: {probe} has no weight_scale -- not NVFP4") +scale2_key = probe.replace(".weight", ".weight_scale_2") +if scale2_key not in index: + sys.exit(f"post-check FAILED: {probe} has no weight_scale_2 -- one-level export") +with safe_open(f"{OUT}/{index[scale_key]}", framework="pt", device="cpu") as f: + s = f.get_tensor(scale_key) + +log(f"post-check {probe}: weight {w.dtype} {tuple(w.shape)} | scale {s.dtype} {tuple(s.shape)}") +assert w.dtype == torch.uint8, f"expected packed uint8 weight, got {w.dtype}" +assert s.dtype == torch.float8_e4m3fn, f"expected float8_e4m3fn scale, got {s.dtype}" + +n_scales = sum(1 for k in index if k.endswith(".weight_scale")) +log(f"post-check OK: {n_scales:,} NVFP4 weight_scale tensors in the checkpoint") +log("DONE") diff --git a/scripts/minimax_m3_nvfp4/run.sh b/scripts/minimax_m3_nvfp4/run.sh new file mode 100755 index 000000000..7104d3af4 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/run.sh @@ -0,0 +1,143 @@ +#!/usr/bin/env bash +set -euo pipefail + +# GPTQ -> NVFP4 quantization of Minimax-M3-0602, calibrated for w4a4. +# +# Quantizes the routed experts of all 57 MoE layers (57 x 128 x 3 = 21,888 +# modules). Attention, shared experts, dense MLP 0-2, routers, norms, embeddings +# and the vision tower pass through untouched -- the format map from +# tore-quant/recipes/minimax-m3-nvfp4. +# +# bash scripts/minimax_m3_nvfp4/run.sh # preflight + full run +# PREFLIGHT_ONLY=1 bash scripts/minimax_m3_nvfp4/run.sh # checks only +# LAYERS=3 bash scripts/minimax_m3_nvfp4/run.sh # single-layer smoke test +# ACT_AMAX=/path/amax_per_layer.json bash .../run.sh # true w4a4 calibration +# +# READ BEFORE RUNNING -- two things this does NOT do: +# +# 1. Without ACT_AMAX this is a *weight-only* calibration. The Hessian is built +# on clean X, so the weights are solved for activations a w4a4 kernel will +# not feed them. See docs/w4a4_hessian.md. Produce a sidecar with +# tore-quant `ptq.py --emit-amax-sidecar`. +# +# 2. The export writes `weight` (packed uint8) + `weight_scale` (float8_e4m3fn) +# + `weight_scale_2` (fp32 per-tensor). It does NOT write `input_scale`; +# a serving stack expecting the full ModelOpt layout needs that added as a +# post-step (derivable from the ACT_AMAX sidecar: amax/6/448 per input site). +# ACT_AMAX changes what the weights are optimized for, not the tensor set. + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "${HERE}/../.." && pwd)" +PY="${PYTHON_BIN:-${REPO}/.venv/bin/python}" + +# /scratch/tonyzhang/models/Minimax-M3-0602 returns EIO on config.json and every +# safetensors shard (btrfs read errors). This copy is intact. +MODEL="${MODEL:-/data/huggingface/Minimax-M3-0602}" +OUT="${OUT:-/data/huggingface/Minimax-M3-0602-NVFP4-GPTQ-w4a4}" +SNAP="${SNAP:-/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/snapshots/27f2a493e10922e878840d27924aabd0bbf5d745}" +OFFLOAD="${OFFLOAD:-/data/sgambhira/gptqmodel-offload/w4a4}" +ACT_AMAX="${ACT_AMAX:-}" +LAYERS="${LAYERS:-}" +NUM_SAMPLES="${NUM_SAMPLES:-64}" +MAX_TOKENS="${MAX_TOKENS:-8192}" +# Empty = per-projection rank thresholds (gate/up 6144, down 3072). Set to a +# number or "N%" to override globally. +FALLBACK_THRESHOLD="${FALLBACK_THRESHOLD:-}" +RANK_MARGIN="${RANK_MARGIN:-1.0}" +LOG="${LOG:-${HERE}/w4a4_run.log}" +PREFLIGHT_ONLY="${PREFLIGHT_ONLY:-0}" + +# Weights alone are ~211 GiB (57 x 3.80 GiB NVFP4 experts) plus ~16.5 GiB of +# passthrough; leave room for the offload dir on the same filesystem. +MIN_FREE_GIB="${MIN_FREE_GIB:-600}" + +say() { printf '\n\033[1m== %s\033[0m\n' "$*"; } +die() { printf '\033[31mFAIL: %s\033[0m\n' "$*" >&2; exit 1; } + +say "Preflight" + +[[ -x "${PY}" ]] || die "python not executable: ${PY} (set PYTHON_BIN)" + +# Read CONTENT, not just stat: /scratch's copy passes `ls` and fails on read. +[[ -d "${MODEL}" ]] || die "model dir missing: ${MODEL}" +head -c 1 "${MODEL}/config.json" >/dev/null 2>&1 || die "cannot READ ${MODEL}/config.json (I/O error?)" +shards=$(ls "${MODEL}"/*.safetensors 2>/dev/null | wc -l) +[[ "${shards}" -gt 0 ]] || die "no safetensors shards in ${MODEL}" +first_shard=$(ls "${MODEL}"/*.safetensors | head -1) +dd if="${first_shard}" of=/dev/null bs=1M count=1 status=none 2>/dev/null \ + || die "cannot READ ${first_shard} -- checkpoint is corrupt" +echo " model ${MODEL} (${shards} shards, readable)" + +[[ -f "${SNAP}/hle-mxfp8-0602-text2158/responses.jsonl" ]] \ + || die "calibration responses.jsonl missing under ${SNAP}" +echo " calib data ${SNAP}" + +if [[ -n "${ACT_AMAX}" ]]; then + [[ -f "${ACT_AMAX}" ]] || die "ACT_AMAX sidecar not found: ${ACT_AMAX}" + "${PY}" -c "import json,sys; d=json.load(open('${ACT_AMAX}')); d=d.get('amax',d); print(f' act amax ${ACT_AMAX} ({len(d)} entries)')" \ + || die "ACT_AMAX is not valid JSON" +else + printf '\033[33m act amax NOT SET -- weight-only calibration, not w4a4-optimal\033[0m\n' +fi + +mkdir -p "${OFFLOAD}" "$(dirname "${OUT}")" +free_gib=$(df -PBG "$(dirname "${OUT}")" | awk 'NR==2{gsub(/G/,"",$4); print $4}') +[[ "${free_gib}" -ge "${MIN_FREE_GIB}" ]] \ + || die "only ${free_gib} GiB free at $(dirname "${OUT}"); need >= ${MIN_FREE_GIB}" +echo " disk ${free_gib} GiB free" + +"${PY}" - <<'PYCHK' || die "no usable CUDA device" +import torch, sys +if not torch.cuda.is_available(): + sys.exit(1) +free, total = torch.cuda.mem_get_info(0) +print(f" gpu {torch.cuda.device_count()}x {torch.cuda.get_device_name(0)}, " + f"{free/2**30:.0f}/{total/2**30:.0f} GiB free on cuda:0") +PYCHK + +if [[ -e "${OUT}" ]]; then + printf '\033[33m WARNING %s already exists; model.save() will overwrite it\033[0m\n' "${OUT}" +fi + +[[ "${PREFLIGHT_ONLY}" == "1" ]] && { say "Preflight only -- stopping"; exit 0; } + +say "Expected RTN fallback rate" +# A module seeing fewer tokens than its own in_features cannot form a full-rank +# Hessian, so it is RTN-quantized rather than GPTQ-solved. Surface the rate up +# front rather than discovering it in the log. +COVERAGE="${HERE}/coverage_layer_expert.json" +if [[ -f "${COVERAGE}" ]]; then + "${PY}" - <5} tok -> RTN {r:,}/{tot:,} = {r/tot:.1%}") +print(" Thresholds are each projection's in_features: H = XtX is (cols x cols) and") +print(" rank(XtX) <= min(n_tokens, cols), so fewer tokens than cols gives a singular H.") +print(" NOTE: more calibration data does NOT lower this. A percentage threshold is") +print(" scale-invariant (RTN iff share < pct/(100*top_k)); only the threshold moves it.") +PYRTN +else + echo " (no coverage_layer_expert.json; run p1b_expert_coverage.py to predict)" +fi + +say "Quantizing" +echo " out ${OUT}" +echo " log ${LOG}" +[[ -n "${LAYERS}" ]] && echo " layers ${LAYERS} (subset)" + +MODEL="${MODEL}" OUT="${OUT}" SNAP="${SNAP}" OFFLOAD="${OFFLOAD}" \ +ACT_AMAX="${ACT_AMAX}" LAYERS="${LAYERS}" NUM_SAMPLES="${NUM_SAMPLES}" \ +MAX_TOKENS="${MAX_TOKENS}" FALLBACK_THRESHOLD="${FALLBACK_THRESHOLD}" \ +RANK_MARGIN="${RANK_MARGIN}" \ +PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" \ + "${PY}" -u "${HERE}/quantize_w4a4.py" 2>&1 | tee "${LOG}" + +say "Done" +echo " checkpoint ${OUT}" +echo " next validate + add input_scale/weight_scale_2 if your serving" +echo " stack expects the ModelOpt NVFP4 layout (see header note 2)" From 7bedb42808faf4df1d2192b27fe23edb482fc6a7 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 18:16:51 -0700 Subject: [PATCH 07/14] fix(nvfp4): config hardening, NaN-loss RTN fallback, and test-placement repair Config (from_quant_config path): accept the `fp4` format alias when the method resolves to NVFP4, matching the constructor path's _normalize_nvfp4_kwargs, and route NVFP4Config targets through the same kwargs normalization. Config (GPTQ validation): format=nvfp4 now rejects what cannot land on the NVFP4 grid -- non-RTN fallback strategies (asymmetric integer grids whose zero-points the packer cannot represent), fallback smoothing (per-block factors make the stored scale non-E4M3-exact, so packing stops being lossless), and mock_quantization (integer rounding, one scale per 128-column block). GPTQ NaN-loss path: under format=nvfp4 a NaN loss now goes to the RTN fallback (which quantizes through NVFP4Quantizer, staying on-grid) instead of the mock-quantization retry, which is off-grid for the same reasons as above. Also repairs tests/test_nvfp4_export.py from commit 64667668, where zero-context partial staging misplaced two lines -- the two-level docstring landed inside a constructor argument list (SyntaxError: positional argument follows keyword argument) and a weight_scale_2 assert landed in the wrong test. The worktree version was always correct (37 tests green); only the committed copy was broken. Co-Authored-By: Claude Fable 5 --- gptqmodel/quantization/config.py | 70 +++++++++++++++++++++++--------- gptqmodel/quantization/gptq.py | 6 +++ tests/test_nvfp4_export.py | 20 +++++++-- 3 files changed, 73 insertions(+), 23 deletions(-) diff --git a/gptqmodel/quantization/config.py b/gptqmodel/quantization/config.py index 288d552ed..bbba75566 100644 --- a/gptqmodel/quantization/config.py +++ b/gptqmodel/quantization/config.py @@ -2954,7 +2954,17 @@ def from_quant_config(cls, quantize_cfg, format: str = None): }: normalized[key] = val else: - normalized[key] = _normalize_format(val) + raw_method = quantize_cfg.get(METHOD_FIELD_CODE, quantize_cfg.get(QUANT_METHOD_FIELD)) + try: + method_hint = _normalize_quant_method(raw_method) if raw_method is not None else None + except ValueError: + method_hint = None + if method_hint == METHOD.NVFP4: + # NVFP4 payloads accept the `fp4` format alias, same as the + # constructor path (`_normalize_nvfp4_kwargs`). + normalized[key] = FORMAT(_normalize_nvfp4_format(val)) + else: + normalized[key] = _normalize_format(val) elif key in field_names: normalized[key] = val else: @@ -3038,6 +3048,8 @@ def from_quant_config(cls, quantize_cfg, format: str = None): normalized = _normalize_gguf_kwargs(normalized) elif target_cls is FP8Config: normalized = _normalize_fp8_kwargs(normalized) + elif target_cls is NVFP4Config: + normalized = _normalize_nvfp4_kwargs(normalized) elif target_cls is BitsAndBytesConfig: normalized = _normalize_bitsandbytes_kwargs(normalized) @@ -3265,18 +3277,6 @@ class GPTQConfig(PreProcessorConfig): default=NVFP4_WEIGHT_BLOCK_SIZE, metadata={"help": "Activation quantization block size, along the last dim."}, ) - - def allowed_quant_methods(self) -> Tuple[METHOD, ...]: - return (METHOD.GPTQ,) - - def supported_export_formats(self) -> Tuple[FORMAT, ...]: - return GPTQ_EXPORT_FORMATS - - def default_desc_act(self) -> bool: - return False - - def __post_init__(self): - desc_act_user_value = self.desc_act act_amax_path: Optional[str] = field( default=None, metadata={ @@ -3295,6 +3295,18 @@ def __post_init__(self): "`layer{N}.moe_input` -- because they share one activation scale." }, ) + + def allowed_quant_methods(self) -> Tuple[METHOD, ...]: + return (METHOD.GPTQ,) + + def supported_export_formats(self) -> Tuple[FORMAT, ...]: + return GPTQ_EXPORT_FORMATS + + def default_desc_act(self) -> bool: + return False + + def __post_init__(self): + desc_act_user_value = self.desc_act act_group_aware_user_value = self.act_group_aware super().__post_init__() @@ -3339,6 +3351,26 @@ def __post_init__(self): self.act_group_aware = False if float(self.mse or 0.0) > 0.0: raise ValueError("QuantizeConfig: `format=nvfp4` does not support `mse` scale search.") + if self.fallback is not None: + # Only the RTN fallback routes through NVFP4Quantizer; every other + # strategy quantizes on an asymmetric integer grid whose zero-points + # the NVFP4 packer cannot represent. + if self.fallback.strategy != FallbackStrategy.RTN: + raise ValueError( + "QuantizeConfig: `format=nvfp4` requires `fallback.strategy=rtn`." + ) + # Smoothing multiplies the block scale by a per-block factor, so the + # stored scale is no longer E4M3-exact and packing stops being lossless. + if self.fallback.smooth is not None: + raise ValueError( + "QuantizeConfig: `format=nvfp4` does not support `fallback.smooth`." + ) + if self.mock_quantization: + # The mock loop rounds on an integer grid and applies one scale per + # 128-column block, neither of which lands on the NVFP4 grid. + raise ValueError( + "QuantizeConfig: `format=nvfp4` does not support `mock_quantization=True`." + ) if self.act_format is not None: # w4a4: the Hessian is built from Q(X), so the activation grid has to be one @@ -3358,6 +3390,12 @@ def __post_init__(self): raise ValueError( "QuantizeConfig: `act_format` does not support `mock_quantization=True`." ) + elif self.act_amax_path is not None or self.act_amax_key_rules is not None: + # A sidecar with no activation quantization silently does nothing; that is + # far more likely a misconfiguration than an intent. + raise ValueError( + "QuantizeConfig: `act_amax_path`/`act_amax_key_rules` require `act_format` to be set." + ) def _resolve_activation_ordering( self, @@ -3390,12 +3428,6 @@ def _update_meta_payload(self, meta_payload: Dict[str, Any]) -> None: "alpha": self.gptaq.alpha, "device": device if isinstance(device, str) else str(device), } - elif self.act_amax_path is not None or self.act_amax_key_rules is not None: - # A sidecar with no activation quantization silently does nothing; that is - # far more likely a misconfiguration than an intent. - raise ValueError( - "QuantizeConfig: `act_amax_path`/`act_amax_key_rules` require `act_format` to be set." - ) else: device = self.foem.device meta_payload["foem"] = { diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index 97e984cc7..1e340af32 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -1355,6 +1355,12 @@ def quantize( if math.isnan(avg_loss): print("Losses sum item:", torch.sum(Losses).item()) if fallback_configured: + if getattr(self.qcfg, "format", None) == FORMAT.NVFP4: + # The mock retry quantizes on an integer grid and reuses one + # scale per block, both off the NVFP4 grid; the RTN fallback + # goes through NVFP4Quantizer and stays on it. + log.info(f"Quantization: Failed due to `NaN` loss for `{self.name}`, using `{resolved_strategy.value}` fallback quantization for `{self.name}`") + return self._fallback_quantize(resolved_strategy, blocksize) log.info(f"Quantization: Failed due to `NaN` loss for `{self.name}`, use mock quantization retry for `{self.name}`") self.qcfg.mock_quantization = True return self.quantize(blocksize=blocksize) diff --git a/tests/test_nvfp4_export.py b/tests/test_nvfp4_export.py index 6d534ecd1..0eeac484b 100644 --- a/tests/test_nvfp4_export.py +++ b/tests/test_nvfp4_export.py @@ -92,6 +92,12 @@ def test_method_nvfp4_resolves_to_weight_only_config(): assert isinstance(shorthand, NVFP4Config) +def test_from_quant_config_accepts_fp4_alias(): + cfg = QuantizeConfig.from_quant_config({"quant_method": "nvfp4", "format": "fp4", "bits": 4}) + assert isinstance(cfg, NVFP4Config) + assert cfg.format == FORMAT.NVFP4 + + @pytest.mark.parametrize( "kwargs", [ @@ -100,6 +106,12 @@ def test_method_nvfp4_resolves_to_weight_only_config(): {"bits": 4, "group_size": 16, "sym": False}, {"bits": 4, "group_size": 16, "desc_act": True}, {"bits": 4, "group_size": 16, "act_group_aware": True}, + # non-RTN fallback strategies quantize off the NVFP4 grid + {"bits": 4, "group_size": 16, "fallback": {"strategy": "mean"}}, + # fallback smoothing breaks E4M3-exact scales, so packing stops being lossless + {"bits": 4, "group_size": 16, "fallback": {"strategy": "rtn", "smooth": "mse"}}, + # the mock loop rounds on an integer grid, off the NVFP4 grid + {"bits": 4, "group_size": 16, "mock_quantization": True}, ], ) def test_format_nvfp4_rejects_incompatible_gptq_settings(kwargs): @@ -145,6 +157,9 @@ def test_nvfp4_quantizer_matches_torchao_grid(): def test_direct_pack_matches_torchao(): + """RTN pack must match torchao's TWO-level path: per-tensor scale amax/(6*448) + plus E4M3 block scales, matching the ModelOpt/Quark/tore-quant layout.""" + torch.manual_seed(1) lin = nn.Linear(128, 64, bias=True) @@ -157,9 +172,6 @@ def test_direct_pack_matches_torchao(): out_features=64, bias=True, register_buffers=True, - """RTN pack must match torchao's TWO-level path: per-tensor scale amax/(6*448) - plus E4M3 block scales, matching the ModelOpt/Quark/tore-quant layout.""" - ) module.pack_original(linear=lin, scales=None, zeros=None) @@ -168,6 +180,7 @@ def test_direct_pack_matches_torchao(): scales_ref, packed_ref = nvfp4_quantize(w.contiguous(), block_size=BLOCK, per_tensor_scale=gs_ref) assert torch.equal(module.weight, packed_ref.view(torch.uint8)) assert torch.equal(module.weight_scale.view(torch.uint8), scales_ref.view(torch.uint8)) + assert torch.equal(module.weight_scale_2, gs_ref) deq = module.dequantize_weight(dtype=torch.float32).T ref = NVFP4Tensor( @@ -183,7 +196,6 @@ def test_gptq_scale_pack_is_lossless(): torch.manual_seed(2) IN, OUT = 128, 64 qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") - assert torch.equal(module.weight_scale_2, gs_ref) quantizer = NVFP4Quantizer(qcfg=qcfg) quantizer.configure(perchannel=True) From 41b8b1770766eedd2f00b2d36fea744e1ad1bd6d Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 18:17:03 -0700 Subject: [PATCH 08/14] chore(scripts): p1b coverage run hardening from the calibration sessions - Point MODEL at /data/huggingface/Minimax-M3-0602: the /scratch copy returns EIO on config.json and every safetensors shard (btrfs read errors) while passing `ls`, so the failure only shows up on read. - WINDOW=5 (was 1): the co-tenant pipeline that forced single-layer windows is gone; 5-layer windows amortize materialize/free and the peak (~100 GiB transient sdpa masks) fits an idle B200. - Checkpoint after every window (CKPT_EVERY=1) so a kill resumes at the last 5-layer boundary; two earlier runs were lost to co-tenant OOM kills. - OOM retry helper for materialization under memory pressure. Co-Authored-By: Claude Fable 5 --- .../minimax_m3_nvfp4/p1b_expert_coverage.py | 61 ++++++++++++++++--- 1 file changed, 51 insertions(+), 10 deletions(-) diff --git a/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py b/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py index 28aee9c19..9c8a01c78 100644 --- a/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py +++ b/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py @@ -11,13 +11,17 @@ """ import json +import os import random +import time import torch import torch.nn as nn from transformers import AutoTokenizer -MODEL = "/data/sgambhira/models/Minimax-M3-0602" +# /scratch/tonyzhang/models/Minimax-M3-0602 now returns EIO on config.json and +# every safetensors shard (btrfs read errors); this copy is intact. +MODEL = "/data/huggingface/Minimax-M3-0602" SNAP = ( "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745" @@ -26,10 +30,29 @@ SEED = 42 NUM_SAMPLES = 64 MAX_TOKENS = 8192 +# A recurring co-tenant pipeline used to claim GPUs 0-3 for ~150GB in bursts, +# which forced WINDOW=1 (~23GB peak). With the GPU idle, WINDOW=5 amortizes the +# per-window materialize/free cost over 5 layers instead of 1. WINDOW = 5 NUM_EXPERTS = 128 -# GPU 0 is occupied by the Phase-1 pilot process during its save; use an idle GPU. -DEVICE = torch.device("cuda:4") +DEVICE = torch.device("cuda:1") +# Co-tenant bursts OOM-killed two previous runs; checkpoint the CPU-side hidden +# states + counts so a restart resumes instead of redoing completed layers. +CKPT = "/data/sgambhira/gptqmodel-offload/coverage_resume.pt" +CKPT_EVERY = 1 # windows between checkpoint saves +OOM_RETRIES = 20 +OOM_WAIT_S = 90 + + +def _retry_oom(fn, what, max_tries=OOM_RETRIES, wait_s=OOM_WAIT_S): + for i in range(max_tries): + try: + return fn() + except torch.OutOfMemoryError: + torch.cuda.empty_cache() + print(f"COV OOM during {what} (attempt {i + 1}/{max_tries}); waiting {wait_s}s", flush=True) + time.sleep(wait_s) + raise RuntimeError(f"COV {what} failed after {max_tries} OOM retries") from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 @@ -90,19 +113,27 @@ def pre_hook(module, args, kwargs): # --- windowed streaming forward --- # hidden states enter each window via inputs_embeds; rotary/causal-mask are # position-derived so per-window model forwards are exact. -embed = model.shell_module_materialize(lm.embed_tokens, DEVICE) model.shell_module_materialize(lm.rotary_emb, DEVICE) model.shell_module_materialize(lm.norm, DEVICE) -hidden = [] -with torch.inference_mode(): - for ids in inputs: - hidden.append(embed(ids.to(DEVICE)).to("cpu")) -print("COV embedded all samples", flush=True) +start_layer = 0 +if os.path.exists(CKPT): + state = torch.load(CKPT, map_location="cpu") + hidden = state["hidden"] + counts.update(state["counts"]) + start_layer = state["next_layer"] + print(f"COV resumed from checkpoint: continuing at layer {start_layer}", flush=True) +else: + embed = model.shell_module_materialize(lm.embed_tokens, DEVICE) + hidden = [] + with torch.inference_mode(): + for ids in inputs: + hidden.append(embed(ids.to(DEVICE)).to("cpu")) + print("COV embedded all samples", flush=True) orig_layers = lm.layers -for w_start in range(0, num_layers, WINDOW): +for w_start in range(start_layer, num_layers, WINDOW): w_layers = all_layers[w_start : w_start + WINDOW] # materialization resolves module paths against the model tree, so the full # layer list must be attached while materializing; truncate only for forward. @@ -130,6 +161,13 @@ def tail_hook(module, args, output): torch.cuda.empty_cache() print(f"COV window done: layers {w_start}-{w_start + len(w_layers) - 1}", flush=True) + next_layer = w_start + len(w_layers) + if ((w_start - start_layer) // WINDOW + 1) % CKPT_EVERY == 0 and next_layer < num_layers: + tmp = CKPT + ".tmp" + torch.save({"hidden": hidden, "counts": counts, "next_layer": next_layer}, tmp) + os.replace(tmp, CKPT) + print(f"COV checkpoint saved through layer {next_layer - 1}", flush=True) + lm.layers = orig_layers for h in hook_handles: @@ -150,6 +188,9 @@ def tail_hook(module, args, output): json.dump(payload, f) print(f"COV wrote {OUT_JSON}", flush=True) +if os.path.exists(CKPT): + os.remove(CKPT) # a finished run must not seed the next fresh run + threshold = 0.005 * total_tokens print("COV layer | experts>0 | experts>=fallback_thr | top1_share | top8_share", flush=True) for li in moe_layers: From 8323b9a2f4968d62d849df1eb8aa21778210b5df Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 18:41:32 -0700 Subject: [PATCH 09/14] fix(dequant): route uint8 E8M0 scales in dequantize_fp8 through block dequant The GPTQ quantize flow does not read weights through the LazyTurtle materialization path fixed in 3e01e1c8. The looper materializes quant targets with role="quant_source" (module_looper.py:1099) and forward modules with role="forward" (:1067); both pull raw tensors via checkpoint_tensors_for_submodule and decode fp8 through _build_decoder_quant_source_module -> dequantize_fp8. That function treated scale_inv as a float multiplier: mxfp8's weight_scale_inv is uint8 UE8M0 exponent bytes (~110-130), which exceed 1, so _fast_scale_arg flipped into divide mode -- dividing by the raw byte value instead of multiplying by 2^(b-127). Measured on Minimax-M3-0602 layers.3 q_proj: absmax 3.94 (divide by ~120) vs correct 0.2812. Every weight GPTQ would have solved against was wrong by a block-dependent factor, silently. A uint8 scale can only be E8M0 exponent bytes -- a real float scale is never serialized as uint8 -- so dequantize_f8_e4m3 now routes uint8 scale/scale_inv straight through dequantize_block_fp8 (verified bit-identical to modelopt's MXFP8QTensor.dequantize). Decoding the bytes in place and continuing was not enough: the fast CPU path also mishandles 2D block-grid scales on small tensors (the reference path, taken for large tensors, is fine), which is why the earlier large-tensor verification passed while a small synthetic failed. Float scales keep their existing paths untouched. Fixes the quant_source and forward decoder roles, TorchFP8Linear, and every other dequantize_fp8 caller in one place. Verified bit-identical to the reference decode on real M3 attention and expert tensors; regression tests pin the multiplicative semantics (4.0 * 2^-3 = 0.5, not 4.0/124), uint8-vs- reference equality, float-path stability, and the real-checkpoint magnitude. Co-Authored-By: Claude Fable 5 --- gptqmodel/quantization/dtype.py | 12 +++++ tests/test_fp8_e8m0_dequant.py | 89 +++++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+) create mode 100644 tests/test_fp8_e8m0_dequant.py diff --git a/gptqmodel/quantization/dtype.py b/gptqmodel/quantization/dtype.py index 4dec74a57..6ce2fd641 100644 --- a/gptqmodel/quantization/dtype.py +++ b/gptqmodel/quantization/dtype.py @@ -616,6 +616,18 @@ def dequantize_f8_e4m3( if scale is not None and scale_inv is not None: raise ValueError("Provide either scale or scale_inv, not both") + + # MX checkpoints (e.g. mxfp8's UE8M0 `weight_scale_inv`) store block scales as + # uint8 exponent bytes: the multiplier is 2^(b - 127). A genuine float scale is + # never serialized as uint8, so route uint8 scales through the block dequant + # helper unconditionally. Feeding the raw bytes onward is wrong twice over: + # magnitudes ~110-130 exceed 1, so `_fast_scale_arg` flips scale_inv into + # divide mode (observed absmax 3.9 vs correct 0.28 on Minimax-M3-0602), and + # the fast CPU path does not understand 2D block grids in the first place. + mx_scale = scale if scale is not None else scale_inv + if mx_scale is not None and mx_scale.dtype == torch.uint8: + return dequantize_block_fp8(tensor, mx_scale, target_dtype=target_dtype) + if tensor.dtype in _FLOAT8_DTYPES and _can_use_fast_path( tensor, scale if scale is not None else scale_inv, diff --git a/tests/test_fp8_e8m0_dequant.py b/tests/test_fp8_e8m0_dequant.py new file mode 100644 index 000000000..0fc55ad76 --- /dev/null +++ b/tests/test_fp8_e8m0_dequant.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +"""uint8 scales entering dequantize_fp8 are UE8M0 exponent bytes, not magnitudes. + +MX checkpoints (mxfp8) serialize `weight_scale_inv` as uint8 E8M0: multiplier +2^(b - 127). Feeding the raw bytes into the float scale arithmetic is the bug +this pins: byte values ~110-130 exceed 1, so `_fast_scale_arg` flipped scale_inv +into divide mode and every dequantized weight came out wrong by a block-dependent +factor (observed absmax 3.9 instead of 0.28 on Minimax-M3-0602). This was the +weight source for GPTQ's quant_source path, so it would have silently poisoned +whole quantization runs. +""" + +import os + + +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + +import pytest +import torch + +from gptqmodel.quantization.dtype import dequantize_block_fp8, dequantize_fp8 + + +BLOCK = 32 # mxfp8 weight_block_size [1, 32] + + +def _mx_tensor(rows=8, cols=128, seed=0): + """A synthetic mxfp8 pair: fp8 payload + uint8 E8M0 block scales.""" + + torch.manual_seed(seed) + payload = (torch.randn(rows, cols) * 100).to(torch.float8_e4m3fn) + # exponents spanning well below and above the bias, like real checkpoints + scale = torch.randint(110, 135, (rows, cols // BLOCK), dtype=torch.uint8) + return payload, scale + + +def test_uint8_scale_inv_decodes_as_e8m0(): + payload, scale = _mx_tensor() + out = dequantize_fp8(payload, scale_inv=scale, axis=None, target_dtype=torch.bfloat16) + ref = dequantize_block_fp8(payload, scale, target_dtype=torch.bfloat16) + assert torch.equal(out, ref), "uint8 scale_inv not decoded as 2^(b-127)" + + +def test_uint8_decode_is_multiplicative_not_divisive(): + """The historic failure mode: bytes > 1 flipped the path into divide mode.""" + + payload = torch.full((1, BLOCK), 4.0).to(torch.float8_e4m3fn) + scale = torch.tensor([[127 - 3]], dtype=torch.uint8) # 2^-3 = 0.125 + out = dequantize_fp8(payload, scale_inv=scale, axis=None, target_dtype=torch.float32) + assert torch.allclose(out, torch.full((1, BLOCK), 0.5)), ( + f"expected 4.0 * 2^-3 = 0.5, got {out[0, 0].item()} " + "(4.0 / 124 would indicate divide-by-raw-bytes has returned)" + ) + + +def test_float_scale_inv_behavior_unchanged(): + """Float scales must not take the uint8/E8M0 branch: their multiplier + semantics are pre-existing and pinned here with a full-shape scale, which is + unambiguous on every internal path.""" + + payload, _ = _mx_tensor(seed=1) + scale_f32 = torch.full(payload.shape, 0.125, dtype=torch.float32) + out = dequantize_fp8(payload, scale_inv=scale_f32, axis=None, target_dtype=torch.float32) + ref = payload.to(torch.float32) * 0.125 + assert torch.allclose(out, ref) + + +@pytest.mark.skipif( + not os.path.isdir("/data/huggingface/Minimax-M3-0602"), reason="M3 checkpoint not present" +) +def test_real_minimax_m3_tensor_decodes_to_sane_magnitude(): + """End-to-end pin against the real checkpoint that exposed the bug.""" + + import json + + from safetensors import safe_open + + d = "/data/huggingface/Minimax-M3-0602" + idx = json.load(open(f"{d}/model.safetensors.index.json"))["weight_map"] + name = "language_model.model.layers.3.self_attn.q_proj.weight" + with safe_open(os.path.join(d, idx[name]), framework="pt") as f: + w = f.get_tensor(name) + with safe_open(os.path.join(d, idx[name + "_scale_inv"]), framework="pt") as f: + s = f.get_tensor(name + "_scale_inv") + out = dequantize_fp8(w, scale_inv=s, axis=None, target_dtype=torch.bfloat16) + absmax = out.float().abs().max().item() + assert 0.1 < absmax < 1.0, f"absmax {absmax}: ~0.28 is correct, 3.9 is divide-mode, 448 is raw" From 805a04e08531c23a3d144e8eadde6ff28c6ef612 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 18:53:11 -0700 Subject: [PATCH 10/14] fix(gptq): resolve activation amax keys against NamedModule.full_name The looper names modules layer-relative -- GPTQ's `self.name` for an M3 expert is `mlp.experts.1.gate_proj`, with no layer index for the sidecar key rules to capture. Resolution now uses NamedModule.full_name (the full dotted path, `...layers.3.mlp.experts.1.gate_proj`), falling back to `self.name` when no NamedModule wraps the module (bare-module tests, HF-optimum path). Found by the LAYERS=3 smoke run: every rule missed, the lookup fell through to the raw name, and the fail-loud KeyError stopped the run before any wrong scale was applied -- the failure mode the guard exists for. The regression test reproduces the exact shape: layer-relative name + full_name on the wrapper, asserting the rule captures the layer index from the latter. Co-Authored-By: Claude Fable 5 --- gptqmodel/quantization/gptq.py | 14 ++++++++++---- tests/test_w4a4_hessian.py | 16 ++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index 1e340af32..6920404ce 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -317,17 +317,23 @@ def _resolve_act_global_scale(self) -> Optional[torch.Tensor]: table = _load_act_amax_table(path) - key = self.name + # `self.name` is layer-relative (`mlp.experts.1.gate_proj`) -- it carries no + # layer index for the key rules to capture. NamedModule.full_name is the + # full dotted path (`...layers.3.mlp.experts.1.gate_proj`), so resolve + # against that when available. + name = getattr(self._named_module, "full_name", None) or self.name + + key = name for pattern, replacement in (getattr(self.qcfg, "act_amax_key_rules", None) or []): - if re.search(pattern, self.name): - key = re.sub(pattern, replacement, self.name) + if re.search(pattern, name): + key = re.sub(pattern, replacement, name) break if key not in table: # Failing loud beats silently falling back to block-only scales, which # would quantize against a different grid than the kernel uses. raise KeyError( - f"GPTQ: activation amax for module '{self.name}' (key '{key}') not found in " + f"GPTQ: activation amax for module '{name}' (key '{key}') not found in " f"'{path}'. Add the entry or extend `act_amax_key_rules`." ) diff --git a/tests/test_w4a4_hessian.py b/tests/test_w4a4_hessian.py index c9367f0f5..ddeb119a6 100644 --- a/tests/test_w4a4_hessian.py +++ b/tests/test_w4a4_hessian.py @@ -224,6 +224,22 @@ def test_invalid_amax_rejected(tmp_path, bad): gptq._resolve_act_global_scale() +def test_key_rules_resolve_against_named_module_full_name(tmp_path): + """The looper names modules layer-relative (`mlp.experts.1.gate_proj`); + only NamedModule.full_name carries the layer index the key rules capture. + Resolving against `self.name` fails on every real looper run.""" + + from types import SimpleNamespace + + path = _sidecar(tmp_path, {"layer3.moe_input": 6.0}) + rules = [(r".*\.layers\.(\d+)\.mlp\.experts\.\d+\.(?:gate_proj|up_proj)$", r"layer\1.moe_input")] + gptq = _gptq_named("mlp.experts.1.gate_proj", _qcfg_sidecar(path, rules)) + gptq._named_module = SimpleNamespace( + full_name="model.language_model.layers.3.mlp.experts.1.gate_proj" + ) + assert gptq._resolve_act_global_scale().item() == pytest.approx(6.0 / 6.0 / 448.0) + + def test_no_sidecar_leaves_global_scale_unset(): gptq = _gptq_named("anything", _qcfg(act=True)) assert gptq._resolve_act_global_scale() is None From 6a9dda83ab4e3527ef9f971315faa3c47dae3f95 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 23:08:49 -0700 Subject: [PATCH 11/14] feat(scripts): source-format splice phase + smoke-test fixes for the w4a4 pipeline splice_source_passthrough.py (new, run.sh phase 3, SPLICE=1 default): GPTQModel's save materializes every module, so passthrough tensors come out dequantized-BF16 under runtime naming -- numerically right but not the serving format map, and not loadable by a source-arch stack (both checkpoints declare MiniMaxM3SparseForConditionalGeneration, but gptqmodel writes model.language_model...mlp.experts.N.gate_proj where the source and serving loaders expect language_model.model...block_sparse_moe.experts.N.w1). The splice walks the SOURCE tensor list and swaps in the quantized experts under source naming: passthrough stays byte-identical MXFP8/BF16/F32, quantized experts gain {weight U8, weight_scale E4M3, weight_scale_2 F32, input_scale F32}, with input_scale = amax/6/448 from the same sidecar the Hessian was calibrated on -- so the serving activation scale and the calibration scale agree by construction. Verified on the LAYERS=3 smoke checkpoint: index keyset exact (46,606 keys), 17/17 sampled passthrough tensors byte-identical, 18/18 sampled expert quads bit-equal with exact f32 input_scale. quantize_w4a4.py: post-check now probes inside the quantized layer set -- the layer-blind probe hit passthrough layer 10 first and failed a good checkpoint. gptq.py: RTN-fallback WARN logs full_name -- self.name is layer-relative and repeats across all 57 layers, which would make multi-layer fallback rosters collide silently. run.sh: COLUMNS=360 so headless logbar stops truncating the per-module loss column to `...`; emits /quant_roster.json mapping every RTN module with observed tokens + threshold (GPTQ-solved = the complement); phase 3 wiring. Co-Authored-By: Claude Fable 5 --- gptqmodel/quantization/gptq.py | 5 +- scripts/minimax_m3_nvfp4/quantize_w4a4.py | 12 +- scripts/minimax_m3_nvfp4/run.sh | 22 +++ .../splice_source_passthrough.py | 180 ++++++++++++++++++ 4 files changed, 216 insertions(+), 3 deletions(-) create mode 100644 scripts/minimax_m3_nvfp4/splice_source_passthrough.py diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index 6920404ce..e7b6a9aa0 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -1065,8 +1065,11 @@ def quantize( use_hessian = False threshold_text = str(getattr(self.fallback, "threshold", None)) threshold_info = f", threshold_raw={threshold_raw}" if threshold_raw is not None and is_percent else "" + # full_name where available: `self.name` is layer-relative and repeats + # across layers, which makes fallback logs ambiguous in multi-layer runs. + fallback_log_name = getattr(self._named_module, "full_name", None) or self.name log.warn( - f"Quantization: Module `{self.name}` -> " + f"Quantization: Module `{fallback_log_name}` -> " f"Using `{resolved_strategy.value}` fallback quantization (observed {self.nsamples} samples, threshold={threshold_text}{threshold_info}, max_total={self.expected_nsamples})." ) self.H = self.create_H(target_device=target_device) diff --git a/scripts/minimax_m3_nvfp4/quantize_w4a4.py b/scripts/minimax_m3_nvfp4/quantize_w4a4.py index fe98ec88d..19a88a802 100644 --- a/scripts/minimax_m3_nvfp4/quantize_w4a4.py +++ b/scripts/minimax_m3_nvfp4/quantize_w4a4.py @@ -176,9 +176,17 @@ def log(msg: str) -> None: from safetensors import safe_open # noqa: E402 index = json.load(open(f"{OUT}/model.safetensors.index.json"))["weight_map"] -probe = next((k for k in index if ".mlp.experts.0.down_proj.weight" in k), None) +# Probe INSIDE the quantized layer set: with LAYERS set, most layers are +# passthrough and legitimately lack NVFP4 scales, so a layer-blind probe fails +# on a perfectly good checkpoint (it did: it hit passthrough layer 10 first). +probe_layer = LAYERS.partition("-")[0] if LAYERS else "3" +probe = next( + (k for k in index + if k.endswith(f".layers.{probe_layer}.mlp.experts.0.down_proj.weight")), + None, +) if probe is None: - sys.exit("post-check FAILED: no routed-expert weight found in the saved index") + sys.exit(f"post-check FAILED: no routed-expert weight for layer {probe_layer} in the saved index") with safe_open(f"{OUT}/{index[probe]}", framework="pt", device="cpu") as f: w = f.get_tensor(probe) diff --git a/scripts/minimax_m3_nvfp4/run.sh b/scripts/minimax_m3_nvfp4/run.sh index 7104d3af4..63c2d576d 100755 --- a/scripts/minimax_m3_nvfp4/run.sh +++ b/scripts/minimax_m3_nvfp4/run.sh @@ -135,9 +135,31 @@ ACT_AMAX="${ACT_AMAX}" LAYERS="${LAYERS}" NUM_SAMPLES="${NUM_SAMPLES}" \ MAX_TOKENS="${MAX_TOKENS}" FALLBACK_THRESHOLD="${FALLBACK_THRESHOLD}" \ RANK_MARGIN="${RANK_MARGIN}" \ PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" \ +COLUMNS="${COLUMNS:-360}" \ "${PY}" -u "${HERE}/quantize_w4a4.py" 2>&1 | tee "${LOG}" +# Phase 3: splice quantized experts into the SOURCE checkpoint so passthrough +# modules keep their original serialization (MXFP8/BF16/F32, source naming) and +# experts gain input_scale from the amax sidecar -- the serving format map. +SPLICE="${SPLICE:-1}" +OUT_FINAL="${OUT_FINAL:-${OUT}-final}" +if [[ "${SPLICE}" == "1" ]]; then + say "Splicing source passthrough -> ${OUT_FINAL}" + SRC="${MODEL}" QNT="${OUT}" OUT="${OUT_FINAL}" ACT_AMAX="${ACT_AMAX}" "${PY}" -u "${HERE}/splice_source_passthrough.py" 2>&1 | tee -a "${LOG}" +fi + say "Done" +# Roster: which modules were RTN'd (named in WARN lines) vs GPTQ-solved (the rest). +ROSTER="${OUT}/quant_roster.json" +"${PY}" - "${LOG}" "${ROSTER}" <<'PYROSTER' || true +import json, re, sys +log, out = sys.argv[1], sys.argv[2] +rtn = {} +for m in re.finditer(r"Module `([\w.]+)` -> Using `rtn` fallback quantization \(observed (\d+) samples, threshold=(\d+)", open(log, errors="replace").read()): + rtn[m.group(1)] = {"observed_tokens": int(m.group(2)), "threshold": int(m.group(3))} +json.dump({"rtn_modules": rtn, "note": "GPTQ-solved = quantized modules not listed here"}, open(out, "w"), indent=1) +print(f" roster {out} ({len(rtn)} RTN modules)") +PYROSTER echo " checkpoint ${OUT}" echo " next validate + add input_scale/weight_scale_2 if your serving" echo " stack expects the ModelOpt NVFP4 layout (see header note 2)" diff --git a/scripts/minimax_m3_nvfp4/splice_source_passthrough.py b/scripts/minimax_m3_nvfp4/splice_source_passthrough.py new file mode 100644 index 000000000..10feb1604 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/splice_source_passthrough.py @@ -0,0 +1,180 @@ +"""Splice GPTQ-NVFP4 experts into the source checkpoint (passthrough = source bytes). + +GPTQModel's save materializes every module, so passthrough tensors come out as +dequantized BF16 -- numerically right, but not the serving format map: attention, +index projections, shared experts, dense MLP 0-2 and non-quantized expert layers +must stay byte-identical MXFP8/BF16/F32 from the source release +(tore-quant/recipes/minimax-m3-nvfp4 format map). + +This builds the final checkpoint the way tore-quant's build_nvfp4_mxfp8_mix does: +walk the SOURCE tensor list, and only for routed experts of quantized layers swap +in the GPTQ output's NVFP4 tensors (renamed to source naming), plus the +`input_scale` the ModelOpt layout expects, derived from the activation amax +sidecar (amax / 6 / 448 per input site -- the same numbers the Hessian was +calibrated with). + + SRC=/data/huggingface/Minimax-M3-0602 \ + QNT=/data/huggingface/Minimax-M3-0602-NVFP4-w4a4-smoke-l3 \ + OUT=/data/huggingface/Minimax-M3-0602-NVFP4-w4a4-smoke-l3-final \ + ACT_AMAX=scripts/minimax_m3_nvfp4/amax_per_layer.json \ + python scripts/minimax_m3_nvfp4/splice_source_passthrough.py +""" + +import json +import os +import re +import shutil +import time + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +SRC = os.environ.get("SRC", "/data/huggingface/Minimax-M3-0602") +QNT = os.environ["QNT"] +OUT = os.environ["OUT"] +ACT_AMAX = os.environ.get("ACT_AMAX", "").strip() +SHARD_BYTES = int(os.environ.get("SHARD_BYTES", str(8 << 30))) + +t0 = time.time() + + +def log(msg: str) -> None: + print(f"[splice +{time.time() - t0:7.1f}s] {msg}", flush=True) + + +# --- naming: source <-> gptqmodel-save --------------------------------------- +# language_model.model.layers.N.block_sparse_moe.experts.E.{w1,w2,w3} +# <-> model.language_model.layers.N.mlp.experts.E.{gate_proj,down_proj,up_proj} +PROJ_TO_QNT = {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"} +SRC_EXPERT_RE = re.compile( + r"^language_model\.model\.layers\.(\d+)\.block_sparse_moe\.experts\.(\d+)\.(w[123])\." +) + + +def qnt_base(layer: str, expert: str, proj: str) -> str: + return f"model.language_model.layers.{layer}.mlp.experts.{expert}.{PROJ_TO_QNT[proj]}" + + +src_index = json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"] +qnt_index = json.load(open(f"{QNT}/model.safetensors.index.json"))["weight_map"] + +# Quantized modules are exactly those the GPTQ output gave a weight_scale_2. +quantized_layers = sorted({ + int(m.group(1)) + for k in qnt_index if k.endswith(".weight_scale_2") + for m in [re.search(r"\.layers\.(\d+)\.mlp\.experts\.", k)] if m +}) +log(f"quantized layers in {QNT}: {quantized_layers}") +if not quantized_layers: + raise SystemExit("no quantized expert layers found in QNT -- nothing to splice") + +amax = {} +if ACT_AMAX: + payload = json.load(open(ACT_AMAX)) + amax = payload.get("amax", payload) + log(f"input_scale source: {ACT_AMAX} ({len(amax)} entries)") +else: + log("WARNING: no ACT_AMAX -- input_scale tensors will NOT be written") + +# --- plan the output tensor list, in source order ----------------------------- +# Per quantized expert projection: drop source {weight, weight_scale_inv}; emit +# {weight, weight_scale, weight_scale_2[, input_scale]} from QNT under source names. +PROJ_SITE = {"w1": "moe_input", "w3": "moe_input", "w2": "w2_input"} + +plan = [] # (out_name, origin, origin_key) origin in {"src", "qnt", "amax"} +emitted = set() +for name in src_index: + m = SRC_EXPERT_RE.match(name) + if m and int(m.group(1)) in quantized_layers: + layer, expert, proj = m.group(1), m.group(2), m.group(3) + base_src = f"language_model.model.layers.{layer}.block_sparse_moe.experts.{expert}.{proj}" + if base_src in emitted: + continue # second source tensor (weight_scale_inv) of the same module + emitted.add(base_src) + base_q = qnt_base(layer, expert, proj) + for suffix in (".weight", ".weight_scale", ".weight_scale_2"): + if base_q + suffix not in qnt_index: + raise SystemExit(f"QNT missing {base_q}{suffix} -- incomplete quantized module") + plan.append((base_src + suffix, "qnt", base_q + suffix)) + if amax: + key = f"layer{layer}.{PROJ_SITE[proj]}" + if key not in amax: + raise SystemExit(f"amax sidecar missing {key}") + plan.append((base_src + ".input_scale", "amax", key)) + else: + plan.append((name, "src", name)) + +log(f"plan: {len(plan)} tensors ({sum(1 for _, o, _ in plan if o != 'src')} spliced, " + f"{sum(1 for _, o, _ in plan if o == 'src')} source passthrough)") + +# --- stream tensors into fresh shards ----------------------------------------- +os.makedirs(OUT, exist_ok=True) +handles = {} + + +def read_tensor(origin: str, key: str) -> torch.Tensor: + if origin == "amax": + return torch.tensor(float(amax[key]) / 6.0 / 448.0, dtype=torch.float32) + root, index = (SRC, src_index) if origin == "src" else (QNT, qnt_index) + shard = index[key] + cache_key = (root, shard) + if cache_key not in handles: + handles.clear() # one open shard per side at a time; source order keeps locality + handles[cache_key] = safe_open(os.path.join(root, shard), framework="pt") + return handles[cache_key].get_tensor(key) + + +weight_map = {} +buffer, buf_bytes, shard_id = {}, 0, 0 +total = 0 + + +def flush(): + global buffer, buf_bytes, shard_id + if not buffer: + return + shard_id += 1 + fname = f"model-{shard_id:05d}.safetensors" + save_file(buffer, os.path.join(OUT, fname)) + for k in buffer: + weight_map[k] = fname + buffer, buf_bytes = {}, 0 + + +for out_name, origin, key in plan: + t = read_tensor(origin, key) + buffer[out_name] = t + buf_bytes += t.numel() * t.element_size() + total += 1 + if buf_bytes >= SHARD_BYTES: + flush() + if shard_id % 8 == 0: + log(f" wrote shard {shard_id} ({total}/{len(plan)} tensors)") +flush() +log(f"wrote {shard_id} shards, {len(weight_map)} tensors") + +# rename shards to the canonical model-XXXXX-of-YYYYY pattern +final_names = {} +for i in range(1, shard_id + 1): + old = f"model-{i:05d}.safetensors" + new = f"model-{i:05d}-of-{shard_id:05d}.safetensors" + os.replace(os.path.join(OUT, old), os.path.join(OUT, new)) + final_names[old] = new +weight_map = {k: final_names[v] for k, v in weight_map.items()} + +with open(f"{OUT}/model.safetensors.index.json", "w") as f: + json.dump({"metadata": {"total_size": sum( + os.path.getsize(os.path.join(OUT, s)) for s in set(weight_map.values()) + )}, "weight_map": weight_map}, f, indent=1) + +# --- config + aux files: source's, with quantized experts removed from the +# mxfp8 ignore/exclude story. Serving-side quantization_config for the mixed +# NVFP4+MXFP8 layout follows the MiniMax-M3-NVFP4-mxfp8skip14 recipe and is left +# to the serving prep step; here we copy source aux files verbatim. +for fname in os.listdir(SRC): + if fname.endswith((".json", ".py", ".jinja", ".txt")) and fname != "model.safetensors.index.json": + shutil.copy2(os.path.join(SRC, fname), os.path.join(OUT, fname)) +log("copied source aux/config files verbatim (serving quantization_config is a separate prep step)") + +log(f"DONE -> {OUT}") From 39ad168a5c7439cc365706362aebb2cdba5b37c9 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Wed, 5 Aug 2026 23:17:01 -0700 Subject: [PATCH 12/14] fix(nvfp4): two-level block scales in the GPTQ loop via a primed global scale The weight_scale_2 fix (64667668) covered the pack stage and the no-scales RTN path, but pack deliberately preserves the loop's scales -- and the loop's NVFP4Quantizer.find_params still derived them single-level: clamp(block_amax/6, 2^-6, 448). For realistic weight magnitudes those absolute values sit below E4M3's smallest normal, so blocks pin at the floor and quantize on a compressed grid. Measured on the LAYERS=3 smoke checkpoint against MiniMax-M3-NVFP4 (ModelOpt RTN of the same weights): 96.8% of block scales at the floor, reconstruction error 1.4-1.8x the donor's (0.12-0.15 vs 0.082 relative) -- on our RTN-fallback modules too, which should have matched the donor's RTN nearly exactly. cos(ours, donor) was 0.993+, so every layout and losslessness check passed while quality was silently off. find_params only ever sees one 16-wide group, so the per-tensor scale must be fixed from the full weight before the loop: prime_global_scale(W), called in the main quantize path and in _fallback_quantize. Primed, find_params derives block scales as E4M3 ratios against it (torchao's two-level path, verified bit-equal); unprimed behaviour is unchanged single-level. The global scale is snapped to a power of two, 2^ceil(log2(amax/(6*448))) (tore-quant's lossless-conversion snap), so the loop's effective scales split exactly at pack time with no plumbing changes: _pack_with_gptq_scales re-derives gs -- or gs/2 when the max block ratio E4M3-rounds down to exactly 224; either decomposition reproduces the identical product, and the in-pack equality guard holds. Cost vs the donor's greedy per-tensor scale is at most one bit of E4M3 range headroom. Tests: primed find_params bit-equal to nvfp4_quantize(per_tensor_scale=gs); floor-clamp fraction must drop below 50% on N(0, 0.02) weights (was ~100%); pack split exact with a power-of-two recovered scale in {gs, gs/2}. Co-Authored-By: Claude Fable 5 --- gptqmodel/quantization/gptq.py | 9 ++++ gptqmodel/quantization/quantizer.py | 47 +++++++++++++++++++-- tests/test_nvfp4_export.py | 65 +++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 4 deletions(-) diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index e7b6a9aa0..d3cee8c85 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -796,6 +796,11 @@ def _fallback_quantize(self, strategy: FallbackStrategy, blocksize: int): scale_chunks = [] zero_chunks = [] + if isinstance(self.quantizer, NVFP4Quantizer): + # Same two-level priming as the main path: the RTN fallback quantizes + # group by group and would otherwise floor-clamp every block scale. + self.quantizer.prime_global_scale(W) + for start in range(0, self.columns, effective_group_size): end = min(start + effective_group_size, self.columns) block = W[:, start:end] @@ -1103,6 +1108,10 @@ def quantize( W = self.module_copy.to(device=self.H.device) del self.module_copy + if isinstance(self.quantizer, NVFP4Quantizer): + # Two-level NVFP4 needs the per-tensor global scale fixed from the FULL + # weight before any per-group find_params -- groups alone floor-clamp. + self.quantizer.prime_global_scale(W) self.quantizer.find_params(W, weight=True) # H = self.H.to(device=self.H.device) diff --git a/gptqmodel/quantization/quantizer.py b/gptqmodel/quantization/quantizer.py index 3ac5f3a55..082554520 100644 --- a/gptqmodel/quantization/quantizer.py +++ b/gptqmodel/quantization/quantizer.py @@ -5,6 +5,8 @@ # adapted from @qwopqwop200 's [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa/tree/cuda), which itself is based on [gptq](https://github.com/IST-DASLab/gptq) +from typing import Optional + import torch import torch.nn as nn @@ -195,6 +197,35 @@ def configure(self, *args, **kwargs): if f32_to_f4_unpacked is None or f4_unpacked_to_f32 is None: raise RuntimeError("NVFP4Quantizer requires torchao with NVFP4 support.") super().configure(*args, **kwargs) + self.global_scale: Optional[torch.Tensor] = None + + def prime_global_scale(self, weight: torch.Tensor) -> torch.Tensor: + """Fix the per-tensor second-level scale from the FULL weight, before the loop. + + NVFP4 is two-level: block scales are E4M3 *relative to* a per-tensor global + scale. Without one, ``find_params`` derives absolute block scales + ``amax/6`` -- for realistic weights those sit below E4M3's smallest normal + (2^-6) and get floor-clamped (measured 96.8% of blocks on Minimax-M3 + experts), compressing the grid and costing ~1.4x reconstruction error. + + The global scale is snapped to a power of two, ``2^ceil(log2(amax/(6*448)))`` + (tore-quant's lossless-conversion snap), for exactness downstream: the + largest block's E4M3 ratio lands in (224, 448] by construction, so + ``_pack_with_gptq_scales`` re-derives exactly this value from the effective + scales and the two-level split reproduces the loop's grid bit-for-bit. + + find_params only ever sees one 16-wide group at a time, so this must be + called once per module with the full weight; unprimed behaviour is the old + single-level derivation. + """ + + amax = weight.detach().abs().amax().to(torch.float32) + if amax > 0: + exponent = torch.ceil(torch.log2(amax / (_F4_E2M1_MAX * _F8_E4M3_MAX))) + self.global_scale = torch.exp2(exponent).reshape(()) + else: + self.global_scale = torch.ones((), dtype=torch.float32) + return self.global_scale def find_params(self, x, weight=False): if not weight: @@ -203,11 +234,19 @@ def find_params(self, x, weight=False): dev = x.device x_flat = x.flatten(1).to(torch.float32) - # torchao nvfp4_quantize: scale = clamp(amax / 6, e4m3_tiny, 448) cast to E4M3. amax = x_flat.abs().amax(dim=1) - scale = amax / _F4_E2M1_MAX - scale = torch.clamp(scale, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) - scale = scale.to(torch.float8_e4m3fn).to(torch.float32) + gs = getattr(self, "global_scale", None) + if gs is not None: + # Two-level (torchao nvfp4_quantize with per_tensor_scale): block ratio + # amax/(6*gs) uses E4M3's full range; effective scale = e4m3 * gs. + ratio = amax / (_F4_E2M1_MAX * gs.to(amax.device)) + ratio = torch.clamp(ratio, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) + scale = ratio.to(torch.float8_e4m3fn).to(torch.float32) * gs.to(amax.device) + else: + # Single-level (torchao without per_tensor_scale): scale = clamp(amax/6, ...). + scale = amax / _F4_E2M1_MAX + scale = torch.clamp(scale, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) + scale = scale.to(torch.float8_e4m3fn).to(torch.float32) shape = [-1] + [1] * (len(x.shape) - 1) self.scale = scale.reshape(shape).to(device=dev) diff --git a/tests/test_nvfp4_export.py b/tests/test_nvfp4_export.py index 0eeac484b..4d085d46e 100644 --- a/tests/test_nvfp4_export.py +++ b/tests/test_nvfp4_export.py @@ -388,3 +388,68 @@ def test_weight_only_nvfp4_end_to_end(tmp_path: Path): with torch.inference_mode(): out = reloaded.model(input_ids=encoded["input_ids"]) assert torch.isfinite(out.logits).all() + + +def test_primed_quantizer_matches_torchao_two_level_grid(): + """With a primed global scale, find_params must land on torchao's two-level + grid -- the fix for 96.8% of block scales floor-clamping at E4M3's 2^-6 on + realistic weight magnitudes (1.4x reconstruction error vs the donor RTN).""" + + torch.manual_seed(3) + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + W = torch.randn(32, 64, dtype=torch.float32) * 0.02 + gs = quantizer.prime_global_scale(W) + assert float(gs) > 0 and (torch.log2(gs) == torch.log2(gs).round()), "global scale must be a power of two" + + # collect the loop's per-group effective scales + ours = [] + for g in range(0, 64, BLOCK): + quantizer.find_params(W[:, g : g + BLOCK], weight=True) + ours.append(quantizer.scale.flatten()) + ours = torch.stack(ours, dim=1) + + ref_scales, _ = nvfp4_quantize(W.contiguous(), block_size=BLOCK, per_tensor_scale=gs.reshape(())) + ref = ref_scales.to(torch.float32) * gs + assert torch.equal(ours, ref), "primed find_params diverged from torchao two-level scales" + + # the floor-clamp regression itself: on these magnitudes the one-level path + # pins ~100% of blocks at 2^-6; two-level must not. + floored = (ours == torch.finfo(torch.float8_e4m3fn).tiny).float().mean() + assert floored < 0.5, f"{floored:.0%} of scales at the E4M3 floor -- still single-level" + + +def test_primed_scales_pack_losslessly_with_recovered_global_scale(): + """pack must re-derive exactly the primed global scale from the effective + scales, so the two-level split reproduces the loop's grid bit-for-bit.""" + + torch.manual_seed(4) + IN, OUT = 128, 32 + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + W = torch.randn(OUT, IN, dtype=torch.float32) * 0.02 + gs = quantizer.prime_global_scale(W) + scales = [] + for g in range(0, IN, BLOCK): + quantizer.find_params(W[:, g : g + BLOCK], weight=True) + scales.append(quantizer.scale.reshape(-1, 1)) + scales = torch.cat(scales, dim=1) + + module = TorchNVFP4Linear( + bits=4, group_size=BLOCK, sym=True, desc_act=False, + in_features=IN, out_features=OUT, bias=False, + ) + packed, block_scale, ws2 = module._pack_with_gptq_scales(W, scales) + # The invariant is exactness of the split, not equality of the decomposition: + # if the max block ratio E4M3-rounds down to exactly 224, pack legitimately + # recovers gs/2 -- the product is identical either way. + assert torch.equal(block_scale.to(torch.float32) * ws2, scales), "two-level split not exact" + lg = torch.log2(ws2.reshape(())) + assert lg == lg.round(), "recovered global scale must stay a power of two" + assert float(lg) in (float(torch.log2(gs)), float(torch.log2(gs)) - 1.0), ( + "recovered global scale should be gs or gs/2 (boundary rounding)" + ) From e8c5760f9c30cfe74420dad772d0e10482bb711e Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Thu, 6 Aug 2026 09:58:09 -0700 Subject: [PATCH 13/14] fix(looper): dynamic per-module fallback overrides must win over loop() baseline base.py passes quantize_config.fallback into module_looper.loop(); it threads through run_layer_stage -> preprocess -> clone_gptq_config_for_module, which applied it AFTER the dynamic-override block -- so every per-module dynamic `fallback` was silently clobbered by the global value in the standard quantize() path. Found via the LAYERS=3 smoke roster: all three expert projections fell back at threshold=6144 even though the per-projection rank thresholds set down_proj to 3072 (its in_features), needlessly RTN'ing 13 down_proj experts whose Hessians were full-rank. The caller-supplied fallback is now the baseline, applied before the dynamic block, so a dynamic override merges onto it and wins. Regression test pins both directions: an overridden module resolves to its dynamic threshold, a non-overridden module keeps the baseline. Co-Authored-By: Claude Fable 5 --- gptqmodel/looper/gptq_processor.py | 7 ++++++- tests/test_w4a4_hessian.py | 24 ++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/gptqmodel/looper/gptq_processor.py b/gptqmodel/looper/gptq_processor.py index a1dc0e7a4..2c08ce1e1 100644 --- a/gptqmodel/looper/gptq_processor.py +++ b/gptqmodel/looper/gptq_processor.py @@ -57,6 +57,12 @@ def clone_gptq_config_for_module( qcfg_clone = copy.deepcopy(qcfg) + # The caller-supplied fallback (base.py passes quantize_config.fallback into + # loop()) is the BASELINE, applied before dynamic so per-module `fallback` + # overrides win. Applying it after (the previous order) silently clobbered + # every dynamic fallback override in the standard quantize() path. + qcfg_clone.fallback = normalize_fallback(fallback, qcfg_clone.fallback) + # dynamic overrides if qcfg.dynamic is not None: qcfg_clone.bits = qcfg.dynamic_get(module_full_name, "bits", qcfg_clone.bits) @@ -102,7 +108,6 @@ def clone_gptq_config_for_module( qcfg_clone._resolve_activation_ordering(desc_act_override, act_group_aware_override) - qcfg_clone.fallback = normalize_fallback(fallback, qcfg_clone.fallback) return qcfg_clone class GPTQProcessor(LoopProcessor): diff --git a/tests/test_w4a4_hessian.py b/tests/test_w4a4_hessian.py index ddeb119a6..97ffac8ee 100644 --- a/tests/test_w4a4_hessian.py +++ b/tests/test_w4a4_hessian.py @@ -266,3 +266,27 @@ def hessian(cfg, name="k"): assert not torch.allclose(block_only, with_scale), ( "global scale had no effect on the Hessian — sidecar is not reaching the quantizer" ) + + +def test_dynamic_fallback_override_survives_caller_baseline(): + """base.py threads the global quantize_config.fallback into loop(); it must be + the baseline, not a clobber -- per-module dynamic `fallback` overrides (the + per-projection rank thresholds) previously never took effect.""" + + from gptqmodel.looper.gptq_processor import clone_gptq_config_for_module + from gptqmodel.quantization.config import Fallback, FallbackStrategy + + qcfg = QuantizeConfig( + bits=4, group_size=BLOCK, format="nvfp4", sym=True, desc_act=False, + dynamic={r".*\.down_proj$": {"fallback": {"strategy": "rtn", "threshold": 3072}}}, + fallback=Fallback(strategy=FallbackStrategy.RTN, threshold=6144), + ) + caller_baseline = Fallback(strategy=FallbackStrategy.RTN, threshold=6144) + + clone = clone_gptq_config_for_module( + qcfg, "model.layers.3.mlp.experts.0.down_proj", fallback=caller_baseline) + assert clone.fallback.threshold == 3072, "dynamic override clobbered by caller fallback" + + clone2 = clone_gptq_config_for_module( + qcfg, "model.layers.3.mlp.experts.0.gate_proj", fallback=caller_baseline) + assert clone2.fallback.threshold == 6144, "non-overridden module must keep the baseline" From 6e03ec039d862cdd335b0e60d7b68f4b9b119403 Mon Sep 17 00:00:00 2001 From: Sirish Gambhira Date: Thu, 6 Aug 2026 10:41:50 -0700 Subject: [PATCH 14/14] fix(nvfp4): remove the single-level scale fallback -- unprimed find_params raises The unprimed branch of NVFP4Quantizer.find_params derived absolute amax/6 block scales, which floor-clamp at E4M3's 2^-6 normal minimum on realistic weight magnitudes (measured 96.8% of blocks on Minimax-M3 experts, 1.4x reconstruction error). Keeping it as a silent fallback meant any caller that forgot prime_global_scale would reproduce that bug undetectably -- it cost a full smoke-run cycle plus a donor-checkpoint comparison to find the first time. An unprimed caller is a bug, not a mode: find_params now raises with the priming instruction. Audit of every path that can hold an NVFP4Quantizer, as the precondition for making this fatal: - GPTQ.quantize: primed before the first full-weight find_params; every other site (static_groups precompute, act-order branches, CPU-fallback re-runs, per-group loop) is downstream in the same call. - GPTQ._fallback_quantize: primed before the RTN per-block loop. - GPTAQ / FOEM: subclass GPTQ and inherit the NVFP4 quantizer but run their own quantize() with full-weight find_params -- these were UNPRIMED and would have produced floor-clamped NVFP4 silently; now primed. - rtn.py constructs a plain Quantizer (never NVFP4); the weight-only NVFP4 lifecycle packs via quantize_nvfp4_weight with no quantizer object. The old unprimed-vs-torchao-single-level parity test is repurposed as the raise regression; the lossless-pack test primes as production does; primed two-level parity was already covered. Co-Authored-By: Claude Fable 5 --- gptqmodel/quantization/foem.py | 3 +++ gptqmodel/quantization/gptaq.py | 3 +++ gptqmodel/quantization/quantizer.py | 28 +++++++++++++++------------- tests/test_nvfp4_export.py | 27 +++++++-------------------- 4 files changed, 28 insertions(+), 33 deletions(-) diff --git a/gptqmodel/quantization/foem.py b/gptqmodel/quantization/foem.py index b17002068..7c070282b 100644 --- a/gptqmodel/quantization/foem.py +++ b/gptqmodel/quantization/foem.py @@ -20,6 +20,7 @@ from ..quantization import QuantizeConfig from ..utils.torch import TORCH_GTE_28, torch_compile, torch_sync from .gptq import GPTQ +from .quantizer import NVFP4Quantizer class FOEM(GPTQ): @@ -155,6 +156,8 @@ def quantize( W = self.module_copy self.module_copy = None + if isinstance(self.quantizer, NVFP4Quantizer): + self.quantizer.prime_global_scale(W) self.quantizer.find_params(W, weight=True) H = self.H diff --git a/gptqmodel/quantization/gptaq.py b/gptqmodel/quantization/gptaq.py index 79ef7f0dd..4472f7f65 100644 --- a/gptqmodel/quantization/gptaq.py +++ b/gptqmodel/quantization/gptaq.py @@ -20,6 +20,7 @@ from ..quantization import QuantizeConfig from ..utils.torch import TORCH_GTE_28, torch_compile, torch_sync from .gptq import GPTQ +from .quantizer import NVFP4Quantizer class GPTAQ(GPTQ): @@ -140,6 +141,8 @@ def quantize( W = self.module_copy self.module_copy = None + if isinstance(self.quantizer, NVFP4Quantizer): + self.quantizer.prime_global_scale(W) self.quantizer.find_params(W, weight=True) H = self.H diff --git a/gptqmodel/quantization/quantizer.py b/gptqmodel/quantization/quantizer.py index 082554520..bdfb340de 100644 --- a/gptqmodel/quantization/quantizer.py +++ b/gptqmodel/quantization/quantizer.py @@ -215,8 +215,7 @@ def prime_global_scale(self, weight: torch.Tensor) -> torch.Tensor: scales and the two-level split reproduces the loop's grid bit-for-bit. find_params only ever sees one 16-wide group at a time, so this must be - called once per module with the full weight; unprimed behaviour is the old - single-level derivation. + called once per module with the full weight; an unprimed find_params raises. """ amax = weight.detach().abs().amax().to(torch.float32) @@ -236,17 +235,20 @@ def find_params(self, x, weight=False): amax = x_flat.abs().amax(dim=1) gs = getattr(self, "global_scale", None) - if gs is not None: - # Two-level (torchao nvfp4_quantize with per_tensor_scale): block ratio - # amax/(6*gs) uses E4M3's full range; effective scale = e4m3 * gs. - ratio = amax / (_F4_E2M1_MAX * gs.to(amax.device)) - ratio = torch.clamp(ratio, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) - scale = ratio.to(torch.float8_e4m3fn).to(torch.float32) * gs.to(amax.device) - else: - # Single-level (torchao without per_tensor_scale): scale = clamp(amax/6, ...). - scale = amax / _F4_E2M1_MAX - scale = torch.clamp(scale, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) - scale = scale.to(torch.float8_e4m3fn).to(torch.float32) + if gs is None: + # No single-level fallback: absolute amax/6 block scales sit below + # E4M3's 2^-6 normal floor on realistic weights (measured 96.8% of + # blocks clamped, 1.4x reconstruction error). An unprimed caller is a + # bug, not a mode. + raise RuntimeError( + "NVFP4Quantizer.find_params called without a primed global scale; " + "call prime_global_scale(full_weight) before per-group find_params." + ) + # Two-level (torchao nvfp4_quantize with per_tensor_scale): block ratio + # amax/(6*gs) uses E4M3's full range; effective scale = e4m3 * gs. + ratio = amax / (_F4_E2M1_MAX * gs.to(amax.device)) + ratio = torch.clamp(ratio, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) + scale = ratio.to(torch.float8_e4m3fn).to(torch.float32) * gs.to(amax.device) shape = [-1] + [1] * (len(x.shape) - 1) self.scale = scale.reshape(shape).to(device=dev) diff --git a/tests/test_nvfp4_export.py b/tests/test_nvfp4_export.py index 4d085d46e..599e82ed5 100644 --- a/tests/test_nvfp4_export.py +++ b/tests/test_nvfp4_export.py @@ -124,31 +124,17 @@ def test_format_nvfp4_rejects_incompatible_gptq_settings(kwargs): # --------------------------------------------------------------------------- -def test_nvfp4_quantizer_matches_torchao_grid(): - torch.manual_seed(0) - W = torch.randn(64, 128, dtype=torch.float32) - - scales_ref, packed_ref = nvfp4_quantize(W, block_size=BLOCK) - deq_ref = _nvfp4_dequant(packed_ref, scales_ref) +def test_nvfp4_quantizer_requires_primed_global_scale(): + """Single-level (unprimed) scale derivation was removed: absolute amax/6 + block scales floor-clamp at E4M3's 2^-6 on realistic weights. An unprimed + find_params is a caller bug and must raise, not silently degrade.""" qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") quantizer = NVFP4Quantizer(qcfg=qcfg) quantizer.configure(perchannel=True) - out = torch.empty_like(W) - scales = [] - for g in range(0, W.shape[1], BLOCK): - block = W[:, g : g + BLOCK] - quantizer.find_params(block, weight=True) - out[:, g : g + BLOCK] = quantizer.quantize(block) - scales.append(quantizer.scale) - scales = torch.cat(scales, dim=1) - - assert torch.equal(out, deq_ref) - assert torch.equal( - scales.to(torch.float8_e4m3fn).view(torch.uint8), - scales_ref.view(torch.uint8), - ) + with pytest.raises(RuntimeError, match="prime_global_scale"): + quantizer.find_params(torch.randn(64, BLOCK), weight=True) # --------------------------------------------------------------------------- @@ -200,6 +186,7 @@ def test_gptq_scale_pack_is_lossless(): quantizer.configure(perchannel=True) W = torch.randn(OUT, IN, dtype=torch.float32) + quantizer.prime_global_scale(W) Q = torch.empty_like(W) scales = [] for g in range(0, IN, BLOCK):