Feat/gptq nvfp4

angelslim/compressor/quant/core/config.py

@@ -174,6 +174,37 @@ def __init__(self, config, global_config=None):
                 "dequant_to_bf16": quantization_args.quant_method.get("dequant_to_bf16", False),
                 "actorder": quantization_args.quant_method.get("actorder", True),
             }
+        elif "nvfp4_gptq" in self.quant_algo or "nvfp4_gptaq" in self.quant_algo:
+            # NVFP4 weight format optimized with the GPTQ error-compensation
+            # loop (weight-only). Shares the GPTQ runner with int4_gptq; the
+            # weight_format flag routes the per-format quant primitives. The
+            # NVFP4 micro-scaling block is fixed at 16. act-order may stay on:
+            # the GPTQ runner folds the input permutation into the previous
+            # layer's output channels, so the 16-element block stays physically
+            # contiguous along K even when columns are reordered.
+            self.act_observer = None
+            self.weight_observer = None
+            self.kv_cache_observer = None
+            block_size = (
+                16
+                if quantization_args.quant_method["group_size"] == -1
+                else quantization_args.quant_method["group_size"]
+            )
+            self.quant_algo_info = {
+                # ``group_size`` is read by the GPTQ runner; ``block_size`` by
+                # the NVFP4 packing path. Keep them in sync.
+                "group_size": block_size,
+                "block_size": block_size,
+                "ignore_layers": quantization_args.ignore_layers,
+                "weight_format": "nvfp4",
+                "checkpoint_format": "nvfp4",
+                "dequant_to_bf16": quantization_args.quant_method.get("dequant_to_bf16", False),
+                "actorder": quantization_args.quant_method.get("actorder", True),
+                "share_gate_up_weight_scale_2": quantization_args.quant_method.get(
+                    "share_gate_up_weight_scale_2", True
+                ),
+                "w": f"nvfp4_{weight_quant_method}",
+            }
         elif "nvfp4" in self.quant_algo:
             is_weight_only = "weight_only" in self.quant_algo
             is_dynamic = "dynamic" if "dynamic" in self.quant_algo else "static"

angelslim/compressor/quant/modules/gptq/gptq.py

@@ -24,9 +24,18 @@
 from huggingface_hub import save_torch_state_dict
 from tqdm import tqdm
 
-from .....utils import decide_device_for_distributed, find_layers, print_info
+from .....utils import (
+    decide_device_for_distributed,
+    find_layers,
+    find_parent_layer_and_sub_name,
+    print_info,
+)
 from ...modules.catcher import Catcher
-from ...modules.helper_layer import GPTQQuantLinear
+from ...modules.helper_layer import (
+    GPTQQuantLinear,
+    NVFP4QDQModule,
+    compute_nvfp4_weight_scale_2,
+)
 from .gptaq_module import GPTAQModule
 from .gptq_module import GPTQModule
 
@@ -56,16 +65,30 @@ def __init__(self, model, seq_length=2048, hidden_size=2560, sym=True, actorder=
         self.quant_bits = self.model.quant_config.quant_bit
         self.group_size = self.model.quant_config.quant_algo_info["group_size"]
         self.ignore_layers = self.model.quant_config.quant_algo_info["ignore_layers"]
+        # Weight number format backend: "int4" (default) or "nvfp4".
+        self.weight_format = self.model.quant_config.quant_algo_info.get("weight_format", "int4")
+        self.block_size = self.model.quant_config.quant_algo_info.get(
+            "block_size", self.group_size
+        )
         self.dequant_to_bf16 = bool(
             self.model.quant_config.quant_algo_info.get("dequant_to_bf16", False)
         )
+        # NVFP4 only: share one per-tensor (level-2) weight scale across each
+        # fused gate/up (and q/k/v) group, matching the single per-tensor scale
+        # a fused-GEMM deployment applies. Default on; set False to keep the
+        # legacy per-layer scale_2.
+        self.share_gate_up_weight_scale_2 = bool(
+            self.model.quant_config.quant_algo_info.get("share_gate_up_weight_scale_2", True)
+        )
         self.percdamp = 0.01
         self.sym = sym
         self.actorder = actorder
         self.seq_length = seq_length
         self.hidden_size = hidden_size
         self.dtype = next(iter(self.layers.parameters())).dtype
         self.quantizers = {}
+        # NVFP4: per-tensor (level-2) scale per quantized layer.
+        self.nvfp4_weight_scales_2 = {}
         self.gptq = {}
         self.quant_algo = self.model.quant_config.quant_algo
         self.native_inp_caches = {}
@@ -110,24 +133,45 @@ def _move_layer_to_device(self, layer, device):
         return layer.to(device)
 
     def get_actorder_prev_names(self, name, subset):
+        # Enable act-order for *_proj layers whose input permutation can be
+        # folded into the output channels of the preceding gate_proj/up_proj.
+        # This keeps the (NVFP4 / int4) group blocks physically contiguous
+        # along K even after column reordering, so no g_idx is needed.
         if not name.endswith("down_proj"):
             return []
 
         prefix = name[: -len("down_proj")]
         prev_names = [f"{prefix}gate_proj", f"{prefix}up_proj"]
-        return [prev_name for prev_name in prev_names if prev_name in subset]
+        if any(prev_name not in subset for prev_name in prev_names):
+            return []
+
+        current_layer = subset[name]
+        for prev_name in prev_names:
+            prev_layer = subset[prev_name]
+            if prev_layer.weight.shape[0] != current_layer.weight.shape[1]:
+                return []
+        return prev_names
 
     def reorder_quantizer_output_channels(self, quantizer_name, input_perm):
         if quantizer_name not in self.quantizers:
             return
 
         scale, zero = self.quantizers[quantizer_name]
         scale_perm = input_perm.to(scale.device)
-        zero_perm = input_perm.to(zero.device)
-        self.quantizers[quantizer_name] = (
-            scale.index_select(0, scale_perm),
-            zero.index_select(0, zero_perm),
-        )
+        # Reorder per-output-channel scale rows. For nvfp4 the second element
+        # is the per-tensor (scalar) weight_scale_2, which is invariant to
+        # output-channel permutation and must NOT be index_select'd.
+        if self.weight_format == "nvfp4":
+            self.quantizers[quantizer_name] = (
+                scale.index_select(0, scale_perm),
+                zero,
+            )
+        else:
+            zero_perm = input_perm.to(zero.device)
+            self.quantizers[quantizer_name] = (
+                scale.index_select(0, scale_perm),
+                zero.index_select(0, zero_perm),
+            )
 
     def fold_input_permutation_into_prev_layers(
         self,
@@ -257,11 +301,70 @@ def _filter_distributed_expert_subset(self, layer, subset):
 
         if skipped > 0:
             print_info(
-                "Filter non-local expert layers for GPTQ: "
-                f"local experts [{start}, {end}), skipped {skipped} layer(s)."
+                f"Filter non-local expert layers for GPTQ: local experts "
+                f"[{start}, {end}), skipped {skipped} layer(s)."
             )
         return filtered_subset
 
+    @staticmethod
+    def _nvfp4_scale_2_group_key(name):
+        """Group key for layers that must share one NVFP4 per-tensor scale_2.
+
+        Deployment fuses each expert's gate_proj+up_proj (and attention's
+        q/k/v) into a single GEMM that can apply only one per-tensor (level-2)
+        scale across the fused weight. GPTQ must therefore quantize every member
+        of such a group against the *same* weight_scale_2, otherwise the half of
+        the fused weight whose scale_2 was overwritten at deploy time is
+        dequantized with the wrong per-tensor scale -> a systematic multiplicative
+        bias that occasionally flips a token.
+
+        Returns (group_prefix, member_tag) for fusible layers, else None.
+        """
+        for suffix in ("gate_proj", "up_proj"):
+            if name.endswith(suffix):
+                return name[: -len(suffix)], "gate_up"
+        for suffix in ("q_proj", "k_proj", "v_proj"):
+            if name.endswith(suffix):
+                return name[: -len(suffix)], "qkv"
+        return None
+
+    def _assign_shared_nvfp4_weight_scale_2(self):
+        """Pre-pass: give fused gate/up (and q/k/v) groups one shared scale_2.
+
+        Runs BEFORE fasterquant so GPTQ compensates against the exact grid that
+        deployment will use. The shared scale_2 is derived from the group-wide
+        amax (max of each member's amax). Because the shared amax is >= each
+        member's own amax, every per-block E4M3 scale only shrinks, so no block
+        scale overflows the E4M3 range that single-layer scaling would have hit.
+        """
+        groups = {}
+        for name in self.gptq:
+            if any(ignore in name for ignore in self.ignore_layers):
+                continue
+            key = self._nvfp4_scale_2_group_key(name)
+            if key is None:
+                continue
+            groups.setdefault(key, []).append(name)
+
+        for (prefix, tag), names in groups.items():
+            if len(names) < 2:
+                # Nothing to fuse with; let fasterquant compute scale_2 itself.
+                continue
+            group_amax = None
+            for name in names:
+                # self.gptq[name].w is the cloned fp32 weight before compensation;
+                # GPTQ compensation preserves amax magnitude, so this is the same
+                # amax fasterquant would otherwise use per-layer.
+                w_amax = self.gptq[name].w.abs().amax().to(torch.float32)
+                group_amax = w_amax if group_amax is None else torch.maximum(group_amax, w_amax)
+            shared_scale_2 = compute_nvfp4_weight_scale_2(group_amax)
+            for name in names:
+                self.gptq[name].weight_scale_2 = shared_scale_2.to(self.gptq[name].dev)
+            print_info(
+                f"NVFP4 shared weight_scale_2 for {tag} group '{prefix}*': "
+                f"{shared_scale_2.item():.6g} across {len(names)} layers"
+            )
+
     @torch.no_grad()
     def run(self, dataloader):
         for model_module in self.layers:
@@ -329,7 +432,12 @@ def run(self, dataloader):
                     self.native_inp_caches[name] = []
                     self.gptq[name] = GPTAQModule(subset[name], quant_bits=self.quant_bits)
                 else:
-                    self.gptq[name] = GPTQModule(subset[name], quant_bits=self.quant_bits)
+                    self.gptq[name] = GPTQModule(
+                        subset[name],
+                        quant_bits=self.quant_bits,
+                        weight_format=self.weight_format,
+                        block_size=self.block_size,
+                    )
 
             def pre_process_fwd_hook(layer_name):
                 def tmp(_, inp, out):
@@ -380,6 +488,9 @@ def tmp(_, inp, out):
             for h in handles:
                 h.remove()
 
+            if self.weight_format == "nvfp4" and self.share_gate_up_weight_scale_2:
+                self._assign_shared_nvfp4_weight_scale_2()
+
             for name in self.gptq:
                 if any(ignore in name for ignore in self.ignore_layers):
                     continue
@@ -389,8 +500,8 @@ def tmp(_, inp, out):
                     and self.gptq[name].nsamples == 0
                 ):
                     print_info(
-                        f"Skip {name} because no calibration samples were routed "
-                        "to this local expert layer."
+                        f"Skip {name} because no calibration samples were "
+                        f"routed to this local expert layer."
                     )
                     self.gptq[name].free()
                     continue
@@ -403,10 +514,23 @@ def tmp(_, inp, out):
                     actorder=actorder,
                     sym=self.sym,
                 )
-                self.quantizers[f"{self.layers_block_name}.{i}.{name}"] = (
-                    scale,
-                    zero,
-                )
+                quant_name = f"{self.layers_block_name}.{i}.{name}"
+                if self.weight_format == "nvfp4":
+                    # NVFP4 uses block_size=16 -> 8x more scale groups than
+                    # int4 (gs=128). Keep the per-layer scales on CPU so they
+                    # do not accumulate on GPU across all 80 layers (the cause
+                    # of the layer-16 OOM). They are only needed at convert().
+                    self.quantizers[quant_name] = (
+                        scale.cpu(),
+                        zero.cpu(),
+                    )
+                    # ``zero`` slot carries the per-tensor level-2 scale.
+                    self.nvfp4_weight_scales_2[quant_name] = zero.cpu()
+                else:
+                    self.quantizers[quant_name] = (
+                        scale,
+                        zero,
+                    )
                 self.fold_input_permutation_into_prev_layers(
                     i,
                     subset,
@@ -511,12 +635,56 @@ def _convert_llm(self):
             force_layer_back_to_cpu=True,
         )
 
+    def _convert_nvfp4(self):
+        """Insert NVFP4QDQModule for each quantized layer (real packing).
+
+        GPTQ has already written the compensated, E2M1-snapped weights back to
+        ``layer.weight`` (bf16). NVFP4QDQModule re-quantizes them into the
+        packed modelopt NVFP4 format (packed uint4 + E4M3 block scale + FP32
+        weight_scale_2) using the scales GPTQ produced, so the deployed grid is
+        exactly what GPTQ optimized against.
+        """
+        model = self.model.model
+        model.cpu()
+        print_info("Packing NVFP4 model...")
+        layers = find_layers(model, layers=self.model.observer_layer_classes)
+
+        with tctl.threadpool_limits(limits=1):
+            pbar = tqdm(self.quantizers.keys(), leave=True)
+            for name in pbar:
+                pbar.set_description(f"Packing {name}...", refresh=True)
+                if name not in layers:
+                    continue
+                sub_layer = layers[name].cpu()
+                block_scale_e4m3, _zero = self.quantizers[name]
+                weight_scale_2 = self.nvfp4_weight_scales_2[name]
+
+                qdq_module = NVFP4QDQModule(
+                    weight=sub_layer.weight,
+                    weight_scale=block_scale_e4m3.cpu(),
+                    weight_scale_2=weight_scale_2.cpu(),
+                    bias=sub_layer.bias,
+                    block_size=self.block_size,
+                    input_scale=None,
+                )
+                parent_layer, sub_name = find_parent_layer_and_sub_name(model, name)
+                setattr(parent_layer, sub_name, qdq_module)
+        print_info("NVFP4 model packed.")
+
     def convert(self):
         """
         Saves scales and inserts QDQ modules.
         """
         print_info("Start convert model...")
-        if self.dequant_to_bf16:
+        if self.weight_format == "nvfp4":
+            if self.dequant_to_bf16:
+                print_info(
+                    "dequant_to_bf16=True: skip NVFP4 packing, keep "
+                    "fake-quantized bf16 weights."
+                )
+            else:
+                self._convert_nvfp4()
+        elif self.dequant_to_bf16:
             print_info(
                 "dequant_to_bf16=True: skip int4 packing, keep fake-quantized bf16 weights."
             )
@@ -600,6 +768,14 @@ def forward(self, x):
                 except Exception:
                     self.model.model.config.quantization_config = None
             self.model.model.config.torch_dtype = "bfloat16"
+        elif self.weight_format == "nvfp4":
+            # Packed modelopt-style NVFP4 checkpoint (two-level scaling).
+            self.model.model.config.quantization_config = {
+                "quant_method": "nvfp4",
+                "kv_cache_scheme": None,
+                "group_size": self.block_size,
+                "exclude_modules": self.ignore_layers,
+            }
         else:
             self.model.model.config.quantization_config = {
                 "bits": self.quant_bits,