Add WMSE weight scale search calibration

CHANGELOG.rst

@@ -26,6 +26,7 @@ Changelog
 - 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 the ``wmse`` (weighted MSE) NVFP4 weight-scale-search calibration algorithm (`ScaleSweep <https://arxiv.org/abs/2606.07618>`_, eqs. 12-13). It is ``local_hessian`` with the per-cin-block Hessian replaced by its diagonal — the per-input-channel importance ``Imp_i = ‖X[:, i]‖²`` — so each block minimizes ``Σ_b Imp_b ΔW_b²``. Captured from the same input-quantizer activation stream, it keeps the candidate scale sweep, the FP8 scale sweep (including a diagonal Triton fast path) and layerwise calibration, at ``block_size`` instead of ``block_size²`` storage per block. Available as ``{'method': 'wmse'}`` and via the ``nvfp4_w4a4_weight_wmse`` preset.
 - 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:
 - Add support for ONNX Q/DQ node placement for DLA via the new flag ``--target_dla``.

examples/hf_ptq/hf_ptq.py

@@ -276,6 +276,7 @@ def make_calib_dataloader(
         "nvfp4_experts_only",
         "nvfp4_omlp_only",
         "nvfp4_w4a4_weight_local_hessian",
+        "nvfp4_w4a4_weight_wmse",
         "mxfp8",
     }
 )

examples/llm_eval/lm_eval_hf.py

@@ -89,6 +89,8 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
     auto_quantize_checkpoint = arg_dict.pop("auto_quantize_checkpoint", None)
     calib_batch_size = arg_dict.pop("calib_batch_size", None)
     calib_size = arg_dict.pop("calib_size", 512)
+    calib_dataset = arg_dict.pop("calib_dataset", "cnn_dailymail")
+    calib_seqlen = arg_dict.pop("calib_seqlen", 512)
     compress = arg_dict.pop("compress", False)
 
     # Sparse attention arguments
@@ -120,6 +122,8 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
             tokenizer=model_obj.tokenizer,
             batch_size=calib_batch_size,
             calib_size=calib_size,
+            data=calib_dataset,
+            max_sample_length=calib_seqlen,
             auto_quantize_bits=auto_quantize_bits,
             auto_quantize_method=auto_quantize_method,
             auto_quantize_score_size=auto_quantize_score_size,
@@ -167,6 +171,8 @@ def create_from_arg_string(
     "quant_cfg",
     "calib_batch_size",
     "calib_size",
+    "calib_dataset",
+    "calib_seqlen",
     "auto_quantize_bits",
     "auto_quantize_method",
     "auto_quantize_score_size",
@@ -192,6 +198,18 @@ def _add_modelopt_args(parser):
     parser.add_argument(
         "--calib_size", type=int, help="Calibration size for quantization", default=512
     )
+    parser.add_argument(
+        "--calib_dataset",
+        type=str,
+        default="cnn_dailymail",
+        help="Dataset used for quantization calibration",
+    )
+    parser.add_argument(
+        "--calib_seqlen",
+        type=int,
+        default=512,
+        help="Maximum sequence length for quantization calibration samples",
+    )
     parser.add_argument(
         "--auto_quantize_bits",
         type=float,

examples/llm_eval/quantization_utils.py

@@ -19,6 +19,8 @@
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.config import need_calibration
+from modelopt.torch.quantization.model_calib import max_calibrate
+from modelopt.torch.quantization.nn import StaticBlockScaleQuantizer, TensorQuantizer
 from modelopt.torch.quantization.plugins import register_hf_attentions_on_the_fly
 from modelopt.torch.utils.dataset_utils import (
     create_forward_loop,
@@ -52,6 +54,99 @@
 }
 
 
+def _nvfp4_ablation_config(method, with_activations):
+    weight_quant_type = "dynamic" if method == "awq_lite" else "static"
+    quant_cfg = [
+        *mtq.config._base_disable_all,
+        {
+            "quantizer_name": "*weight_quantizer",
+            "enable": True,
+            "cfg": {
+                "num_bits": (2, 1),
+                "block_sizes": {-1: 16, "type": weight_quant_type, "scale_bits": (4, 3)},
+            },
+        },
+        # Layerwise calibration bootstraps the first decoder-layer inputs with a
+        # full-model forward before static weight quantizers have been calibrated
+        # and promoted. Keep embeddings quantized, but use the normal dynamic
+        # NVFP4 path so that bootstrap remains executable.
+        {
+            "parent_class": "nn.Embedding",
+            "quantizer_name": "*weight_quantizer",
+            "enable": True,
+            "cfg": {
+                "num_bits": (2, 1),
+                "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            },
+        },
+    ]
+    if with_activations:
+        quant_cfg.append(
+            {
+                "quantizer_name": "*input_quantizer",
+                "enable": True,
+                "cfg": {
+                    "num_bits": (2, 1),
+                    "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                },
+            }
+        )
+    quant_cfg.extend(
+        {"quantizer_name": pattern, "enable": False}
+        for pattern in (
+            "*embed_vision*",
+            "*vision_tower*",
+            "*visual*",
+            "*vision_model*",
+            "*multi_modal_projector*",
+        )
+    )
+
+    algorithm = {"method": method, "layerwise": {"enable": True}}
+    if method in ("local_hessian", "mse", "wmse"):
+        algorithm["fp8_scale_sweep"] = True
+    return {"quant_cfg": quant_cfg, "algorithm": algorithm}
+
+
+CUSTOM_CONFIG["ABLATE_W4A4_WMSE"] = _nvfp4_ablation_config("wmse", True)
+CUSTOM_CONFIG["ABLATE_W4A16_WMSE"] = _nvfp4_ablation_config("wmse", False)
+
+
+def _calibrate_layerwise_external_static_weights(model):
+    """Max-calibrate static weights outside decoder layers after layerwise PTQ.
+
+    Layerwise calibration intentionally operates on decoder layers only. Ablation
+    configs re-enable normally skipped modules such as ``lm_head``, so calibrate
+    and promote any static weight quantizers that the decoder pass did not touch.
+    Already promoted decoder weights are left unchanged.
+    """
+    calibrated = []
+    for name, module in model.named_modules():
+        if not name or not hasattr(module, "iter_weights_for_calibration"):
+            continue
+        has_unpromoted_static_weight = False
+        for _, quantizer in module.iter_weights_for_calibration():
+            if (
+                isinstance(quantizer, TensorQuantizer)
+                and not isinstance(quantizer, StaticBlockScaleQuantizer)
+                and quantizer.is_enabled
+                and quantizer.is_static_block_quant
+            ):
+                has_unpromoted_static_weight = True
+                break
+        if not has_unpromoted_static_weight:
+            continue
+        max_calibrate(
+            module,
+            forward_loop=None,
+            distributed_sync=False,
+            shared_states={"weight_global_amax": {"patterns": []}},
+        )
+        calibrated.append(name)
+    if calibrated:
+        print(f"Max-calibrated layerwise-external static weights: {calibrated}")
+
+
 def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN, trust_remote_code=False):
     """Returns the tokenizer from the model ckpt_path."""
     print(f"Initializing tokenizer from {ckpt_path}")
@@ -157,6 +252,8 @@ def forward_step(model, batch):
             register_hf_attentions_on_the_fly(net)
 
         net = mtq.quantize(net, mtq_cfg, calibrate_loop)
+        if isinstance(quant_cfg, str) and quant_cfg.startswith("ABLATE_"):
+            _calibrate_layerwise_external_static_weights(net)
     mtq.print_quant_summary(net)
     # Compress or fold weights for faster evaluation.
     if compress:
@@ -172,6 +269,7 @@ def quantize_model(
     batch_size,
     calib_size,
     data="cnn_dailymail",
+    max_sample_length=512,
     test_generated=True,
     compress=False,
     auto_quantize_bits=None,
@@ -189,6 +287,7 @@ def quantize_model(
         batch_size: the calibration batch size for each calibration inference run.
         calib_size: the total calibration dataset size.
         data: the name of the calibration dataset.
+        max_sample_length: the maximum sequence length of each calibration sample.
         test_generated:  If ``True``, test the generated text before and after quantization.
         compress: If ``True``, compress the model after quantization.
         auto_quantize_bits: The effective bits constraint for auto_quantize.
@@ -224,6 +323,7 @@ def quantize_model(
         tokenizer=tokenizer,
         batch_size=batch_size,
         num_samples=calib_size,
+        max_sample_length=max_sample_length,
         device=device,
         include_labels=is_gradient_based,
     )

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

@@ -184,6 +184,7 @@ def _fp8_scale_sweep_hessian_kernel(
     NUM_CANDIDATES: tl.constexpr,
     ROWS_PER_PROGRAM: tl.constexpr,
     HAS_COUPLING: tl.constexpr,
+    DIAGONAL: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
     cin_block = pid % N_CIN_BLOCKS
@@ -207,12 +208,16 @@ def _fp8_scale_sweep_hessian_kernel(
         ).to(tl.float32)
 
     idx = tl.arange(0, BLOCK_SIZE)
-    hessian = tl.load(
-        hessian_ptr
-        + cin_block * (BLOCK_SIZE * BLOCK_SIZE)
-        + idx[:, None] * BLOCK_SIZE
-        + idx[None, :]
-    ).to(tl.float32)  # [BS, BS]
+    if DIAGONAL:
+        # wmse: only the per-input-channel importance diag(H) is stored, [BS] per cin-block.
+        hessian = tl.load(hessian_ptr + cin_block * BLOCK_SIZE + idx).to(tl.float32)  # [BS]
+    else:
+        hessian = tl.load(
+            hessian_ptr
+            + cin_block * (BLOCK_SIZE * BLOCK_SIZE)
+            + idx[:, None] * BLOCK_SIZE
+            + idx[None, :]
+        ).to(tl.float32)  # [BS, BS]
 
     best_loss = tl.full([ROWS_PER_PROGRAM], float("inf"), dtype=tl.float32)
     best_idx = tl.zeros([ROWS_PER_PROGRAM], dtype=tl.int32)
@@ -223,9 +228,14 @@ def _fp8_scale_sweep_hessian_kernel(
         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 * (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 DIAGONAL:
+            # Σ_b Imp_b Δw_b² — the quadratic form with H replaced by diag(Imp). Squaring
+            # before weighting matches the reference einsum's operand order bit-for-bit.
+            loss = tl.sum((dw * dw) * hessian[None, :], axis=1)  # [ROWS]
+        else:
+            # 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
@@ -248,7 +258,9 @@ def nvfp4_fp8_scale_sweep_hessian(
     Hessian-weighted counterpart of :func:`nvfp4_fp8_scale_sweep`: for each NVFP4 block
     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. When ``coupling_bias`` is
+    :class:`NVFP4MSECalibrator` for ``local_hessian`` calibration. A rank-2 ``hessian`` is
+    read as the diagonal ``diag(Imp)`` and minimizes ``Σ_b Imp_b Δw_b²`` instead (``wmse``).
+    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.
 
@@ -258,7 +270,8 @@ def nvfp4_fp8_scale_sweep_hessian(
             ``b % (cin // block_size)``.
         global_amax: Scalar FP32 global amax (``= reduce_amax(per_block_amax)``).
         hessian: Per-cin-block Hessian of shape ``[cin // block_size, block_size, block_size]``,
-            fp32 (typically normalized by sample count).
+            or its diagonal (per-input-channel importance) of shape
+            ``[cin // 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.
@@ -267,10 +280,14 @@ def nvfp4_fp8_scale_sweep_hessian(
         ``best_amax`` of shape ``[N_BLOCKS]``, fp32, on the same device as ``x``.
     """
     n_blocks, x_flat, best_amax = _prepare_block_sweep(x, block_size)
-    if hessian.dim() != 3 or hessian.shape[1] != block_size or hessian.shape[2] != block_size:
+    diagonal = hessian.dim() == 2
+    if hessian.shape[1:] not in (
+        torch.Size([block_size]),
+        torch.Size([block_size, block_size]),
+    ):
         raise ValueError(
-            f"hessian must have shape [n_cin_blocks, {block_size}, {block_size}], "
-            f"got {tuple(hessian.shape)}."
+            f"hessian must have shape [n_cin_blocks, {block_size}] or "
+            f"[n_cin_blocks, {block_size}, {block_size}], got {tuple(hessian.shape)}."
         )
     n_cin_blocks = hessian.shape[0]
     if n_blocks % n_cin_blocks != 0:
@@ -314,6 +331,7 @@ def nvfp4_fp8_scale_sweep_hessian(
             NUM_CANDIDATES=int(candidate_amaxes.numel()),
             ROWS_PER_PROGRAM=_HESSIAN_ROWS_PER_PROGRAM,
             HAS_COUPLING=coupling_bias is not None,
+            DIAGONAL=diagonal,
             num_warps=_HESSIAN_NUM_WARPS,
         )
     return best_amax

modelopt/torch/quantization/calib/mse.py

@@ -179,9 +179,10 @@ class NVFP4MSECalibrator(MseCalibrator):
     Python sweep. Both fast paths require the input on CUDA in the blocked
     ``[n_blocks, block_size]`` layout with Triton + the kernel package importable:
 
-    - **Hessian-weighted** (local_hessian): taken when ``hessian is not None`` — minimizes
-      ``Δ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.
+    - **Hessian-weighted** (local_hessian / wmse): taken when ``hessian is not None`` — minimizes
+      ``Δwᵀ H Δw`` plus the optional activation error coupling term, or ``Σ_b Imp_b Δw_b²`` when
+      ``hessian`` is the rank-2 diagonal (wmse). 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
@@ -201,9 +202,10 @@ def __init__(
     ):
         """Initialize NVFP4 MSE calibrator with per-block and global amax.
 
-        ``hessian`` (per-cin-block ``[cin // block_size, block_size, block_size]``) enables
-        the Hessian-weighted Triton fast path (local_hessian); ``error_func`` carries the
-        same metric for the reference fallback when the fast path is unavailable.
+        ``hessian`` (per-cin-block ``[cin // block_size, block_size, block_size]``, or its
+        diagonal ``[cin // block_size, block_size]`` for wmse) enables the Hessian-weighted
+        Triton fast path; ``error_func`` carries the same metric for the reference fallback
+        when the fast path is unavailable.
         """
         super().__init__(amax=amax, axis=axis, quant_func=quant_func, error_func=error_func)
         self._global_amax = global_amax.to(dtype=torch.float32)
@@ -253,7 +255,7 @@ def _can_use_triton_fast_path(self, x: torch.Tensor) -> bool:
         return self._error_func is None and self._triton_sweep_eligible(x)
 
     def _can_use_hessian_fast_path(self, x: torch.Tensor) -> bool:
-        """Whether the Hessian-weighted Triton fast path is usable (local_hessian)."""
+        """Whether the Hessian-weighted Triton fast path is usable (local_hessian / wmse)."""
         return self._hessian is not None and self._triton_sweep_eligible(x)
 
     @torch.no_grad()