Add activation-aware local Hessian scale calibration

CHANGELOG.rst

@@ -24,6 +24,7 @@ Changelog
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.
+- Extend ``local_hessian`` block-wise output MSE with an opt-in activation error coupling term, minimizing ``‖X_q·W_q − X·W_0‖²`` per block. The local Hessian now uses input-quantizer outputs whenever input quantization is enabled.
 - Add a fused Triton fast path for ``local_hessian`` NVFP4 weight-scale search (the Hessian-weighted FP8-E4M3 scale sweep). For each NVFP4 block it minimizes ``dwᵀ H dw`` over the 126 candidate scales using the per-cin-block local Hessian on tensor cores, replacing the per-weight Python reference sweep — roughly **34x** faster on a single 8192x4096 weight and bit-exact with the reference for fp32/fp16 weights. Used automatically during ``local_hessian`` calibration for both dense and fused-MoE expert weights; falls back to the reference sweep on CPU, when Triton is unavailable, or via ``MODELOPT_NVFP4_TRITON_SWEEP=0``.
 - Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag.
 - Add Minitron pruning support for Megatron-Core models with the following new attention and MoE variants. For these, only ``hidden_size`` is pruned (alongside the usual ``ffn_hidden_size`` / ``num_layers`` / MoE dimensions); the variant-internal dimensions noted below are not pruned:

examples/hf_ptq/README.md

@@ -162,7 +162,7 @@ export HF_PATH=<the downloaded LLaMA checkpoint from the Hugging Face hub, or si
 scripts/huggingface_example.sh --model $HF_PATH --quant <QFORMAT> --tp [1|2|4|8]
 ```
 
-Supported `QFORMAT` values: `fp8`, `fp8_pc_pt`, `fp8_pb_wo`, `int8`, `int8_sq`, `int8_wo`, `int4_awq`, `w4a8_awq`, `nvfp4`, `nvfp4_awq`, `nvfp4_mse`, `nvfp4_mlp_only`, `nvfp4_experts_only`, `nvfp4_omlp_only`, `nvfp4_svdquant`, `nvfp4_local_hessian`, `w4a8_nvfp4_fp8`, `w4a8_mxfp4_fp8`, `mxfp8`.
+Supported `QFORMAT` values: `fp8`, `fp8_pc_pt`, `fp8_pb_wo`, `int8`, `int8_sq`, `int8_wo`, `int4_awq`, `w4a8_awq`, `nvfp4`, `nvfp4_awq`, `nvfp4_mse`, `nvfp4_mlp_only`, `nvfp4_experts_only`, `nvfp4_omlp_only`, `nvfp4_svdquant`, `nvfp4_local_hessian`, `nvfp4_w4a4_weight_local_hessian_act_error_coupling`, `w4a8_nvfp4_fp8`, `w4a8_mxfp4_fp8`, `mxfp8`.
 
 > *By default `trust_remote_code` is set to false. Please turn it on if model calibration and eval requires it using `--trust_remote_code`.*
 

modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py

@@ -174,6 +174,7 @@ def nvfp4_fp8_scale_sweep(
 def _fp8_scale_sweep_hessian_kernel(
     x_ptr,  # [COUT * N_CIN_BLOCKS * BLOCK_SIZE], any float dtype (loaded as fp32)
     hessian_ptr,  # [N_CIN_BLOCKS * BLOCK_SIZE * BLOCK_SIZE] fp32
+    coupling_bias_ptr,  # same flat layout as x, fp32 when HAS_COUPLING
     candidate_scales_ptr,  # [NUM_CANDIDATES] fp32: per-candidate FP8-quantized block scale
     candidate_amaxes_ptr,  # [NUM_CANDIDATES] fp32: per-candidate block amax (kernel output value)
     best_amax_ptr,  # [COUT * N_CIN_BLOCKS] fp32 output
@@ -182,6 +183,7 @@ def _fp8_scale_sweep_hessian_kernel(
     BLOCK_SIZE: tl.constexpr,
     NUM_CANDIDATES: tl.constexpr,
     ROWS_PER_PROGRAM: tl.constexpr,
+    HAS_COUPLING: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
     cin_block = pid % N_CIN_BLOCKS
@@ -197,6 +199,12 @@ def _fp8_scale_sweep_hessian_kernel(
     ).to(tl.float32)  # [ROWS, BS]
     w_abs = tl.abs(w)
     w_sign = tl.where(w >= 0, 1.0, -1.0)
+    if HAS_COUPLING:
+        coupling_bias = tl.load(
+            coupling_bias_ptr + block_idx[:, None] * BLOCK_SIZE + elem[None, :],
+            mask=row_mask[:, None],
+            other=0.0,
+        ).to(tl.float32)
 
     idx = tl.arange(0, BLOCK_SIZE)
     hessian = tl.load(
@@ -214,10 +222,12 @@ def _fp8_scale_sweep_hessian_kernel(
         scale = tl.load(candidate_scales_ptr + k).to(tl.float32)
         scale_safe = tl.where(scale == 0.0, 1.0, scale)  # scale == 0 only if global_amax == 0
         q_mag = fp4_round_magnitude(w_abs / scale_safe)
-        dw = w_sign * (w_abs - q_mag * scale_safe)  # = w - quant(w), [ROWS, BS]
+        dw = w_sign * (q_mag * scale_safe - w_abs)  # = quant(w) - w, [ROWS, BS]
         # dwᵀ H dw per row (H symmetric); allow_tf32=False keeps it true fp32 vs the reference.
         hdw = tl.dot(dw, hessian, allow_tf32=False)  # [ROWS, BS]
         loss = tl.sum(hdw * dw, axis=1)  # [ROWS]
+        if HAS_COUPLING:
+            loss += 2.0 * tl.sum(dw * coupling_bias, axis=1)
         is_better = loss < best_loss
         best_loss = tl.where(is_better, loss, best_loss)
         best_idx = tl.where(is_better, k, best_idx)
@@ -231,13 +241,16 @@ def nvfp4_fp8_scale_sweep_hessian(
     global_amax: torch.Tensor,
     hessian: torch.Tensor,
     block_size: int = 16,
+    coupling_bias: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Find the per-block FP8 scale minimizing the Hessian-weighted NVFP4 quant error.
 
     Hessian-weighted counterpart of :func:`nvfp4_fp8_scale_sweep`: for each NVFP4 block
-    it minimizes ``dwᵀ H dw`` (``dw = w - quant(w)``) over the 126 FP8 E4M3 candidates,
+    it minimizes ``Δwᵀ H Δw`` (``Δw = quant(w) - w``) over the 126 FP8 E4M3 candidates,
     where ``H`` is the per-cin-block local Hessian shared across all output rows. Used by
-    :class:`NVFP4MSECalibrator` for ``local_hessian`` calibration.
+    :class:`NVFP4MSECalibrator` for ``local_hessian`` calibration. When ``coupling_bias`` is
+    supplied, it adds the activation error coupling ``2 Δwᵀ(PW0)``. The scale-independent
+    activation-error constant is omitted, so candidate losses may be negative.
 
     Args:
         x: Weight tensor on CUDA in the blocked ``[N_BLOCKS, block_size]`` layout, row-major
@@ -247,6 +260,8 @@ def nvfp4_fp8_scale_sweep_hessian(
         hessian: Per-cin-block Hessian of shape ``[cin // block_size, block_size, block_size]``,
             fp32 (typically normalized by sample count).
         block_size: NVFP4 block size (typically 16).
+        coupling_bias: Optional fp32-compatible CUDA tensor with ``x.numel()`` values in the
+            same flat layout as ``x``, containing ``P W0`` for each block.
 
     Returns:
         ``best_amax`` of shape ``[N_BLOCKS]``, fp32, on the same device as ``x``.
@@ -262,6 +277,13 @@ def nvfp4_fp8_scale_sweep_hessian(
         raise ValueError(
             f"n_blocks ({n_blocks}) is not divisible by n_cin_blocks ({n_cin_blocks})."
         )
+    if coupling_bias is not None:
+        if not coupling_bias.is_cuda:
+            raise ValueError("coupling_bias must be a CUDA tensor.")
+        if coupling_bias.numel() != x.numel():
+            raise ValueError(
+                f"coupling_bias must have {x.numel()} elements, got {coupling_bias.numel()}."
+            )
 
     cout = n_blocks // n_cin_blocks
     grid = (triton.cdiv(cout, _HESSIAN_ROWS_PER_PROGRAM) * n_cin_blocks,)
@@ -274,9 +296,15 @@ def nvfp4_fp8_scale_sweep_hessian(
             candidate_amaxes, global_amax_f32, quantize_block_scales=True
         ).to(dtype=torch.float32)
         hessian_flat = hessian.contiguous().to(device=x.device, dtype=torch.float32).view(-1)
+        coupling_bias_flat = (
+            coupling_bias.contiguous().to(device=x.device, dtype=torch.float32).view(-1)
+            if coupling_bias is not None
+            else x_flat
+        )
         _fp8_scale_sweep_hessian_kernel[grid](
             x_flat,
             hessian_flat,
+            coupling_bias_flat,
             candidate_scales,
             candidate_amaxes,
             best_amax,
@@ -285,6 +313,7 @@ def nvfp4_fp8_scale_sweep_hessian(
             BLOCK_SIZE=block_size,
             NUM_CANDIDATES=int(candidate_amaxes.numel()),
             ROWS_PER_PROGRAM=_HESSIAN_ROWS_PER_PROGRAM,
+            HAS_COUPLING=coupling_bias is not None,
             num_warps=_HESSIAN_NUM_WARPS,
         )
     return best_amax

modelopt/torch/quantization/calib/mse.py

@@ -180,7 +180,8 @@ class NVFP4MSECalibrator(MseCalibrator):
     ``[n_blocks, block_size]`` layout with Triton + the kernel package importable:
 
     - **Hessian-weighted** (local_hessian): taken when ``hessian is not None`` — minimizes
-      ``dwᵀ H dw``. Wins over the plain path, so it fires even when ``error_func`` is also set.
+      ``Δwᵀ H Δw`` plus the optional activation error coupling term. Wins over the plain path,
+      so it fires even when ``error_func`` is also set.
     - **plain squared-error**: taken when ``hessian is None and error_func is None``.
 
     Otherwise (CPU, non-blocked layout, Triton unavailable, or an ``error_func`` with no
@@ -196,6 +197,7 @@ def __init__(
         quant_func: Callable | None = None,
         error_func: Callable | None = None,
         hessian: torch.Tensor | None = None,
+        coupling: torch.Tensor | None = None,
     ):
         """Initialize NVFP4 MSE calibrator with per-block and global amax.
 
@@ -206,6 +208,7 @@ def __init__(
         super().__init__(amax=amax, axis=axis, quant_func=quant_func, error_func=error_func)
         self._global_amax = global_amax.to(dtype=torch.float32)
         self._hessian = hessian
+        self._coupling = coupling
         # Set by collect() after either sweep path; consumed by compute_amax.
         self._best_amax: torch.Tensor | None = None
 
@@ -264,8 +267,23 @@ def collect(self, x: torch.Tensor):
         if self._can_use_hessian_fast_path(x):
             from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep_hessian
 
+            hessian = self._hessian
+            assert hessian is not None
+            coupling_bias = None
+            if self._coupling is not None:
+                # Local import avoids model_calib -> calib.mse -> model_calib initialization cycle.
+                from modelopt.torch.quantization.model_calib import _activation_error_coupling_bias
+
+                cout = x.numel() // (hessian.shape[0] * x.shape[-1])
+                coupling_bias = _activation_error_coupling_bias(
+                    x, self._coupling, cout, x.shape[-1]
+                ).reshape(-1)
             best_flat = nvfp4_fp8_scale_sweep_hessian(
-                x.detach(), self._global_amax, self._hessian, block_size=x.shape[-1]
+                x.detach(),
+                self._global_amax,
+                hessian,
+                block_size=x.shape[-1],
+                coupling_bias=coupling_bias,
             )
             self._best_amax = best_flat.reshape(self._initial_amax.shape).to(dtype=torch.float32)
             return

modelopt/torch/quantization/config.py

@@ -1014,9 +1014,9 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
     quantization. It minimizes the output reconstruction error by weighting the loss
     with the local Hessian matrix computed from input activations.
 
-    The local Hessian loss for each block is: ``(dw @ H @ dw.T)`` where:
-    - ``dw = weight - quantized_weight`` (weight reconstruction error per block)
-    - ``H = X @ X.T`` is the local Hessian computed from input activations X
+    The default local Hessian loss is ``ΔWᵀ H ΔW``. The optional activation error coupling
+    extends it to ``ΔWᵀ H ΔW + 2 ΔWᵀ P W0``, where ``ΔW = Wq-W0``,
+    ``H = XqᵀXq / B``, and ``P = Xqᵀ(Xq-X) / B`` for each local cin-block.
 
     """
 
@@ -1061,6 +1061,19 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
         "Default is 16 for NVFP4.",
     )
 
+    activation_error_coupling: bool | None = ModeloptField(
+        default=False,
+        title="Include the activation error coupling term in block-wise output MSE.",
+        description=(
+            "If True, build the local Hessian from fake-quantized activations X_q and add the "
+            "activation error coupling term 2·ΔWᵀ·P·W0, where "
+            "P = (1/B)·X_qᵀ·(X_q - X). This minimizes ‖X_q·W_q - X·W0‖² per block after "
+            "dropping a scale-independent constant, so the relative objective may be negative. "
+            "Layers whose input quantizer is disabled, absent, or uses pre_quant_scale/rotation "
+            "fall back to the Hessian-only objective. Default False preserves existing behavior."
+        ),
+    )
+
     distributed_sync: bool | None = ModeloptField(
         default=True,
         title="Whether to sync the amax across the distributed processes.",
@@ -1669,6 +1682,9 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]:
 NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG: dict[str, Any] = _load_quantize_config_dict(
     "configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian"
 )
+NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_ACT_ERROR_COUPLING_CFG: dict[str, Any] = _load_quantize_config_dict(
+    "configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_error_coupling"
+)
 MAMBA_MOE_NVFP4_AGGRESSIVE_CFG: dict[str, Any] = _load_quantize_config_dict(
     "configs/ptq/presets/model/mamba_moe_nvfp4_aggressive"
 )
@@ -1760,6 +1776,7 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]:
     "MAMBA_MOE_FP8_CONSERVATIVE_CFG",
     "MAMBA_MOE_FP8_AGGRESSIVE_CFG",
     "NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG",
+    "NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_ACT_ERROR_COUPLING_CFG",
     "NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG",
 }