NVFP4 support

docs/w4a4_hessian.md

@@ -0,0 +1,286 @@
+# 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.
+
+## 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 `<base>.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. 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.
+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.
+
+## 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.

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):

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: