diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8e6b14e2468..e6553020ce0 100755 --- a/CHANGELOG.rst +++ b/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 `_, 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``. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 8dcc78afa27..517b5eaab1a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/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", } ) diff --git a/examples/llm_eval/lm_eval_hf.py b/examples/llm_eval/lm_eval_hf.py index 51c0930e8f2..997bf4f5ec9 100755 --- a/examples/llm_eval/lm_eval_hf.py +++ b/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, diff --git a/examples/llm_eval/quantization_utils.py b/examples/llm_eval/quantization_utils.py index 80117bd1627..433a53fb1f0 100644 --- a/examples/llm_eval/quantization_utils.py +++ b/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, ) diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py index fd6cd0d389b..ce9e9f84f59 100644 --- a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py +++ b/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 diff --git a/modelopt/torch/quantization/calib/mse.py b/modelopt/torch/quantization/calib/mse.py index 88dd64b7a36..e60b2cd01ea 100644 --- a/modelopt/torch/quantization/calib/mse.py +++ b/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() diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index c8eb109ca0c..0a892d61391 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -864,7 +864,7 @@ def validate_layerwise_checkpoint_dir(self): class _SharedStatesConfig(ModeloptBaseConfig): - """The ``shared_states`` grouping knob, shared by max / mse / local_hessian calibration.""" + """The ``shared_states`` grouping knob, shared by max / mse / local_hessian / wmse calib.""" shared_states: dict[str, dict[str, list[str]]] | None = ModeloptField( default=None, @@ -1007,21 +1007,14 @@ class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ) -class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): - """Configuration for local Hessian-weighted MSE calibration. - - This algorithm uses activation information to optimize per-block scales for weight - quantization. It minimizes the output reconstruction error by weighting the loss - with the local Hessian matrix computed from input activations. - - 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. +class _ActivationWeightedCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): + """Knobs shared by the activation-weighted weight-scale searches (local_hessian / wmse). + Both build a per-cin-block weighting term from the input activations captured during a + calibration forward and hand it to the MSE weight search; they differ only in whether that + term is the full local Hessian or its diagonal. """ - method: Literal["local_hessian"] = ModeloptField("local_hessian") - step_size: float | None = ModeloptField( default=0.1, gt=0.0, @@ -1061,6 +1054,34 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): "Default is 16 for NVFP4.", ) + distributed_sync: bool | None = ModeloptField( + default=True, + title="Whether to sync the amax across the distributed processes.", + description="If True, the amax will be synced across the distributed processes.", + ) + + debug: bool | None = ModeloptField( + default=False, + title="Debug mode.", + description="If True, module's local Hessian metadata will be kept as a module attribute.", + ) + + +class LocalHessianCalibConfig(_ActivationWeightedCalibConfig): + """Configuration for local Hessian-weighted MSE calibration. + + This algorithm uses activation information to optimize per-block scales for weight + quantization. It minimizes the output reconstruction error by weighting the loss + with the local Hessian matrix computed from input activations. + + 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. + + """ + + method: Literal["local_hessian"] = ModeloptField("local_hessian") + activation_error_coupling: bool | None = ModeloptField( default=False, title="Include the activation error coupling term in block-wise output MSE.", @@ -1074,17 +1095,18 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ), ) - distributed_sync: bool | None = ModeloptField( - default=True, - title="Whether to sync the amax across the distributed processes.", - description="If True, the amax will be synced across the distributed processes.", - ) - debug: bool | None = ModeloptField( - default=False, - title="Debug mode.", - description="If True, module's local Hessian metadata will be kept as a module attribute.", - ) +class WmseCalibConfig(_ActivationWeightedCalibConfig): + """Configuration for weighted-MSE (WMSE) calibration. + + WMSE (`ScaleSweep `_, eqs. 12-13) is the local-Hessian + objective with the per-block Hessian ``H`` replaced by its diagonal — the per-input-channel + importance ``Imp_i = ‖X[:, i]‖²``. Each block therefore minimizes ``Σ_b Imp_b ΔW_b²``, which + keeps the activation weighting while dropping the cross-channel terms (and their ``B x B`` + per-block storage). + """ + + method: Literal["wmse"] = ModeloptField("wmse") class SmoothQuantCalibConfig(QuantizeAlgorithmConfig): @@ -1278,7 +1300,9 @@ def _gptq_qdq_default(self): return self -_ScaleCalibConfig: TypeAlias = MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig +_ScaleCalibConfig: TypeAlias = ( + MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig | WmseCalibConfig +) class LSQConfig(QuantizeAlgorithmConfig): @@ -1347,7 +1371,7 @@ class LSQConfig(QuantizeAlgorithmConfig): default=None, title="Scale calibration algorithm to run first.", description=( - "Dict with 'method' key: 'mse', 'local_hessian', or 'max'. " + "Dict with 'method' key: 'mse', 'local_hessian', 'wmse', or 'max'. " "Optional keys include 'fp8_scale_sweep' for FP4 formats. " "Defaults to {'method': 'mse'} if None." ), @@ -1685,6 +1709,9 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]: 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" ) +NVFP4_W4A4_WEIGHT_WMSE_CFG: dict[str, Any] = _load_quantize_config_dict( + "configs/ptq/presets/model/nvfp4_w4a4_weight_wmse" +) MAMBA_MOE_NVFP4_AGGRESSIVE_CFG: dict[str, Any] = _load_quantize_config_dict( "configs/ptq/presets/model/mamba_moe_nvfp4_aggressive" ) diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 00187d291c0..4016bf14d03 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -113,8 +113,8 @@ def _restore_shared_quant_state_aliases( """Rebuild shared-state ties before checkpoint tensor values are loaded.""" if not metadata.get("shared_quant_states"): return - # max / mse / local_hessian all carry ``shared_states`` (via _SharedStatesConfig) and use - # the same grouping; resolve the patterns that were in effect at save and rebuild the ties. + # max / mse / local_hessian / wmse all carry ``shared_states`` (via _SharedStatesConfig) and + # use the same grouping; resolve the patterns in effect at save and rebuild the ties. patterns = SharedWeightGlobalAmaxState.resolve_patterns( shared_states=getattr(config, "shared_states", None) ) diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index 80be73daa2a..677e3dc1530 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -47,6 +47,7 @@ QuantizeConfig, SmoothQuantCalibConfig, SVDQuantConfig, + WmseCalibConfig, _QuantizeExportConfig, ) from .conversion import ( @@ -68,6 +69,7 @@ mse_calibrate, smoothquant, svdquant, + wmse_calibrate, ) __all__ = ["BaseCalibrateModeDescriptor"] @@ -455,6 +457,22 @@ def config_class(self) -> type[QuantizeAlgorithmConfig]: _calib_func = local_hessian_calibrate +@CalibrateModeRegistry.register_mode +class WmseModeDescriptor(BaseCalibrateModeDescriptor): + """Mode for weighted-MSE (WMSE) calibration algorithm. + + The local-Hessian objective with the per-block Hessian replaced by its diagonal, the + per-input-channel activation importance. + """ + + @property + def config_class(self) -> type[QuantizeAlgorithmConfig]: + """Specifies the config class for the mode.""" + return WmseCalibConfig + + _calib_func = wmse_calibrate + + @CalibrateModeRegistry.register_mode class SmoothQuantModeDescriptor(BaseCalibrateModeDescriptor): """Mode for smoothquant calibration algorithm.""" diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 21fd9429d14..149d8e01ed3 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -68,6 +68,7 @@ "max_calibrate", "smoothquant", "svdquant", + "wmse_calibrate", ] @@ -528,7 +529,8 @@ def _make_weight_mse_calibrator( """Create the MSE calibrator for one eligible weight quantizer (``None`` if ineligible). ``error_func`` overrides the squared-error metric (local-Hessian's per-block weighting). - ``metric`` carries the per-cin-block Hessian and optional activation error coupling matrix, + ``metric`` carries the per-cin-block Hessian (full ``[n, bs, bs]`` for ``local_hessian`` or + diagonal ``[n, bs]`` for ``wmse``) and the optional activation error coupling matrix, enabling NVFP4's Hessian-weighted Triton fast path; ``error_func`` then serves only as the reference fallback. """ @@ -718,12 +720,18 @@ class _LocalHessianAccumulator: scale. An enabled input quantizer gives ``H = ΣX_qᵀX_q``; a disabled one returns its input unchanged and gives ``H = ΣXᵀX``. Buffers are allocated lazily so never-routed experts cost nothing. + + With ``diagonal=True`` (the ``wmse`` algorithm) only the diagonal of that Hessian — the + per-input-channel importance ``Imp_i = ‖X[:, i]‖²`` — is accumulated, so + ``hessian_per_block`` has shape ``[n_blocks, block_size]`` instead of + ``[n_blocks, block_size, block_size]``. Downstream consumers switch on that rank. """ - def __init__(self, cout: int, cin: int, block_size: int): + def __init__(self, cout: int, cin: int, block_size: int, diagonal: bool = False): self.cout = cout self.cin = cin self.block_size = block_size + self.diagonal = diagonal self.num_blocks_per_cin = cin // block_size # Not block-divisible -> no Hessian (falls back to plain MSE). self.is_enabled = cin % block_size == 0 @@ -748,7 +756,7 @@ def accumulate( else None ) x = x.reshape(self.num_blocks_per_cin, self.block_size, -1) - hessian_batch = x @ x.transpose(-1, -2) + hessian_batch = x.square().sum(-1) if self.diagonal else x @ x.transpose(-1, -2) if self.hessian_per_block is None: self.hessian_per_block = hessian_batch else: @@ -763,7 +771,7 @@ def accumulate( self.num_samples += quantizer_output.numel() // self.cin def normalized_hessian(self) -> torch.Tensor | None: - """Per-cin-block Hessian ``H / num_samples`` (``None`` if no samples). + """Per-cin-block Hessian (or diagonal importance) ``H / num_samples``, else ``None``. Shared by both the Triton fast path and the reference ``error_func`` so the two consume one tensor; cached because the accumulated buffer may be freed afterwards. @@ -807,7 +815,12 @@ def local_hessian_error(x: torch.Tensor, xq: torch.Tensor) -> torch.Tensor: original_shape = x.shape # Per-block weighted error: dw (cout,n,bs) · H (n,bs,bs) -> (cout,n). dw = (xq - x).view(cout, -1, bs) - block_loss = torch.einsum("cnb,nbd,cnd->cn", dw, hessian, dw) + block_loss = ( + # wmse: Σ_b Imp_b Δw_b² — the quadratic form with H replaced by diag(Imp). + torch.einsum("cnb,nb->cn", dw * dw, hessian) + if hessian.dim() == 2 + else torch.einsum("cnb,nbd,cnd->cn", dw, hessian, dw) + ) if coupling is not None: bias = _activation_error_coupling_bias(x, coupling, cout, bs) block_loss = block_loss + 2.0 * (dw * bias).sum(-1) @@ -817,18 +830,25 @@ def local_hessian_error(x: torch.Tensor, xq: torch.Tensor) -> torch.Tensor: return local_hessian_error -def _warn_if_block_size_mismatch(weight_quantizer: TensorQuantizer, block_size: int, name: str): +def _warn_if_block_size_mismatch( + weight_quantizer: TensorQuantizer, + block_size: int, + name: str, + algorithm: str = "local_hessian", +): """Warn if the Hessian block_size differs from the quantizer's scale block (misaligns).""" block_sizes = getattr(weight_quantizer, "block_sizes", None) quant_block = block_sizes.get(-1) if block_sizes else None if quant_block is not None and quant_block != block_size: warn_rank_0( - f"local_hessian: block_size ({block_size}) != quantizer scale block " + f"{algorithm}: block_size ({block_size}) != quantizer scale block " f"({quant_block}) for {name}; Hessian weighting will not align with the scale blocks." ) -def _warn_local_hessian_fallback(name, weight, weight_quantizer, block_size, warned: set): +def _warn_local_hessian_fallback( + name, weight, weight_quantizer, block_size, warned: set, algorithm: str = "local_hessian" +): """Warn once per ``(name, cin)`` when a captured layer falls back to plain MSE.""" if weight.dim() < 2: return @@ -838,10 +858,10 @@ def _warn_local_hessian_fallback(name, weight, weight_quantizer, block_size, war warned.add((name, cin)) if cin % block_size != 0: warn_rank_0( - f"local_hessian: {name} input features ({cin}) not divisible by block_size " + f"{algorithm}: {name} input features ({cin}) not divisible by block_size " f"({block_size}); falling back to plain MSE for these weights." ) - _warn_if_block_size_mismatch(weight_quantizer, block_size, name) + _warn_if_block_size_mismatch(weight_quantizer, block_size, name, algorithm) def _is_quant_fused_experts(module: nn.Module) -> bool: @@ -869,6 +889,7 @@ def _register_local_hessian_input_hooks( block_size, warned, activation_error_coupling: bool = False, + algorithm: str = "local_hessian", ): """Register forward hooks feeding each weight's input activations to ``capture``. @@ -917,7 +938,7 @@ def _capture_coupling(input_quantizer, layer_name): if weight is None or weight.dim() != 2 or not module.weight_quantizer.is_enabled: continue _warn_local_hessian_fallback( - name, weight, module.weight_quantizer, block_size, warned + name, weight, module.weight_quantizer, block_size, warned, algorithm ) input_quantizer = getattr(module, "input_quantizer", None) @@ -954,7 +975,12 @@ def _dense_hook( if weight is None or quantizers is None or input_quantizer is None: continue _warn_local_hessian_fallback( - f"{name}.{weight_name}", weight[0], quantizers[0], block_size, warned + f"{name}.{weight_name}", + weight[0], + quantizers[0], + block_size, + warned, + algorithm, ) # Snapshot which experts are enabled now, before the caching forward silences # all weight quantizers — so we don't capture (and discard) disabled experts. @@ -976,63 +1002,34 @@ def _dense_hook( @torch.no_grad() -def local_hessian_calibrate( +def _hessian_weighted_calibrate( model: nn.Module, - forward_loop: ForwardLoop | None = None, - distributed_sync: bool = True, - step_size: float = 0.1, - start_multiplier: float = 0.25, - stop_multiplier: float = 4.0, - fp8_scale_sweep: bool = True, - block_size: int = 16, - activation_error_coupling: bool = False, - debug: bool = False, - shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, + forward_loop: ForwardLoop | None, + distributed_sync: bool, + step_size: float, + start_multiplier: float, + stop_multiplier: float, + fp8_scale_sweep: bool, + block_size: int, + activation_error_coupling: bool, + debug: bool, + shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None, + algorithm: str, + diagonal: bool, ): - """Calibrate weight quantizers by minimizing the Hessian-weighted error. - - Minimizes ``(Wq - W)ᵀ H (Wq - W)`` with a per-block Hessian built from the input - quantizer output. Thus ``H = ΣXᵀX`` when that quantizer is disabled and - ``H = ΣX_qᵀX_q`` when enabled. It is captured with weight fake-quant disabled and fed - to :func:`mse_calibrate`'s weight search via ``error_func``. + """Shared accumulate -> metric -> MSE-search pipeline for ``local_hessian`` and ``wmse``. - Like :func:`mse_calibrate`, TensorQuantizer weights are calibrated — with the Hessian - metric where a weight pairs with its input activations (dense linears and HF fused-MoE - experts), plain MSE otherwise. Other quantizer types (e.g. SequentialQuantizer) are - unsupported and left at their max-calibrated scale. - - With ``activation_error_coupling=True``, each eligible local block additionally uses - ``2 ΔWᵀ P W`` with - ``P = X_qᵀ(X_q-X) / B``. This is the activation error coupling term in the exact - block-wise expansion of ``||X_q W_q - X W||²`` after dropping its scale-independent - constant. The objective can therefore be negative. Both matrices retain only local - cin-blocks; layers without an eligible enabled input quantizer use the Hessian-only path. - - Args: - model: Model to be calibrated. - forward_loop: A callable which takes the model as argument and - forwards calibration data through the model. Required for this algorithm. - distributed_sync: Whether to sync amax across distributed processes. - step_size: Step size for amax search (default: 0.1). - start_multiplier: Starting multiplier for amax search (default: 0.25). - stop_multiplier: Ending multiplier for amax search (default: 4.0). - fp8_scale_sweep: If True, sweep over all 128 possible FP8 E4M3 scale values - for NVFP4 per-block quantization (default: True). - block_size: Block size for local Hessian computation (default: 16). - activation_error_coupling: Include the activation error coupling term in block-wise - output MSE. Default False omits this additional term. - debug: If True, retain the per-quantizer Hessian accumulators on the model - (``model._local_hessian_accumulators``) for inspection. - - See :class:`LocalHessianCalibConfig ` - for details on the configuration options. + ``diagonal`` accumulates only the per-input-channel importance ``Imp_i = ‖X[:, i]‖²`` + (``wmse``) instead of the full per-cin-block Hessian (``local_hessian``); everything else + — the hooks, the max-calibration warm start, the candidate sweep — is identical. + ``algorithm`` only names the method in user-facing messages. """ if forward_loop is None: - warnings.warn("forward_loop must be provided for local_hessian; skipping local_hessian") + warnings.warn(f"forward_loop must be provided for {algorithm}; skipping {algorithm}") return # Phase 1: max-calibrate (also bootstraps dead experts + promotes/syncs NVFP4 static). - print_rank_0("local_hessian: Running max calibration for all quantizers...") + print_rank_0(f"{algorithm}: Running max calibration for all quantizers...") max_calibrate(model, forward_loop, distributed_sync, shared_states=shared_states) name_to_module = dict(model.named_modules()) @@ -1053,7 +1050,9 @@ def capture(weight_quantizer, weight, quantizer_output, unquantized_input=None): ) acc = accumulators.get(id(weight_quantizer)) if acc is None: - acc = _LocalHessianAccumulator(weight.shape[0], weight.shape[1], block_size) + acc = _LocalHessianAccumulator( + weight.shape[0], weight.shape[1], block_size, diagonal=diagonal + ) accumulators[id(weight_quantizer)] = acc acc.accumulate(quantizer_output_local, unquantized_input_local) @@ -1067,8 +1066,9 @@ def capture(weight_quantizer, weight, quantizer_output, unquantized_input=None): block_size, warned, activation_error_coupling=activation_error_coupling, + algorithm=algorithm, ) - print_rank_0("local_hessian: Caching activations and computing local Hessian...") + print_rank_0(f"{algorithm}: Caching activations and computing the weighting term...") try: with set_quantizer_by_cfg_context( model, [{"quantizer_name": "*weight_quantizer", "enable": False}] @@ -1082,8 +1082,8 @@ def capture(weight_quantizer, weight, quantizer_output, unquantized_input=None): # amax sync runs before this), so refined amaxes can diverge. All-reduce Hessian / re-sync. if dist_is_initialized() and dist_size() > 1: warn_rank_0( - "local_hessian: Hessian is not synced across ranks; refined weight amaxes may " - "diverge under tensor/data parallelism. Treat local_hessian as single-rank for now." + f"{algorithm}: the weighting term is not synced across ranks; refined weight amaxes " + f"may diverge under tensor/data parallelism. Treat {algorithm} as single-rank for now." ) # Phase 3: weight search. Build error_funcs first so build_error_func caches the normalized @@ -1097,7 +1097,7 @@ def capture(weight_quantizer, weight, quantizer_output, unquantized_input=None): for qid, acc in accumulators.items() if acc.normalized_hessian() is not None } - print_rank_0("local_hessian: Running MSE calibration with local Hessian loss...") + print_rank_0(f"{algorithm}: Running MSE calibration with the weighted loss...") _mse_calibrate_weights( model, name_to_module, @@ -1126,7 +1126,137 @@ def capture(weight_quantizer, weight, quantizer_output, unquantized_input=None): if torch.cuda.is_available(): torch.cuda.empty_cache() - print_rank_0("local_hessian: Calibration complete.") + print_rank_0(f"{algorithm}: Calibration complete.") + + +@torch.no_grad() +def local_hessian_calibrate( + model: nn.Module, + forward_loop: ForwardLoop | None = None, + distributed_sync: bool = True, + step_size: float = 0.1, + start_multiplier: float = 0.25, + stop_multiplier: float = 4.0, + fp8_scale_sweep: bool = True, + block_size: int = 16, + activation_error_coupling: bool = False, + debug: bool = False, + shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, +): + """Calibrate weight quantizers by minimizing the Hessian-weighted error. + + Minimizes ``(Wq - W)ᵀ H (Wq - W)`` with a per-block Hessian built from the input + quantizer output. Thus ``H = ΣXᵀX`` when that quantizer is disabled and + ``H = ΣX_qᵀX_q`` when enabled. It is captured with weight fake-quant disabled and fed + to :func:`mse_calibrate`'s weight search via ``error_func``. + + Like :func:`mse_calibrate`, TensorQuantizer weights are calibrated — with the Hessian + metric where a weight pairs with its input activations (dense linears and HF fused-MoE + experts), plain MSE otherwise. Other quantizer types (e.g. SequentialQuantizer) are + unsupported and left at their max-calibrated scale. + + With ``activation_error_coupling=True``, each eligible local block additionally uses + ``2 ΔWᵀ P W`` with + ``P = X_qᵀ(X_q-X) / B``. This is the activation error coupling term in the exact + block-wise expansion of ``||X_q W_q - X W||²`` after dropping its scale-independent + constant. The objective can therefore be negative. Both matrices retain only local + cin-blocks; layers without an eligible enabled input quantizer use the Hessian-only path. + + Args: + model: Model to be calibrated. + forward_loop: A callable which takes the model as argument and + forwards calibration data through the model. Required for this algorithm. + distributed_sync: Whether to sync amax across distributed processes. + step_size: Step size for amax search (default: 0.1). + start_multiplier: Starting multiplier for amax search (default: 0.25). + stop_multiplier: Ending multiplier for amax search (default: 4.0). + fp8_scale_sweep: If True, sweep over all 128 possible FP8 E4M3 scale values + for NVFP4 per-block quantization (default: True). + block_size: Block size for local Hessian computation (default: 16). + activation_error_coupling: Include the activation error coupling term in block-wise + output MSE. Default False omits this additional term. + debug: If True, retain the per-quantizer Hessian accumulators on the model + (``model._local_hessian_accumulators``) for inspection. + + See :class:`LocalHessianCalibConfig ` + for details on the configuration options. + """ + _hessian_weighted_calibrate( + model, + forward_loop, + distributed_sync, + step_size, + start_multiplier, + stop_multiplier, + fp8_scale_sweep, + block_size, + activation_error_coupling, + debug, + shared_states, + algorithm="local_hessian", + diagonal=False, + ) + + +@torch.no_grad() +def wmse_calibrate( + model: nn.Module, + forward_loop: ForwardLoop | None = None, + distributed_sync: bool = True, + step_size: float = 0.1, + start_multiplier: float = 0.25, + stop_multiplier: float = 4.0, + fp8_scale_sweep: bool = True, + block_size: int = 16, + debug: bool = False, + shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, +): + """Calibrate weight quantizers by minimizing the activation-weighted squared error. + + Weighted MSE (`ScaleSweep `_, eqs. 12-13) is + :func:`local_hessian_calibrate` with the per-block Hessian replaced by its diagonal: the + per-input-channel importance ``Imp_i = ‖X[:, i]‖²``, so each block minimizes + ``Σ_b Imp_b ΔW_b²``. ``Imp`` is captured from the same input-quantizer outputs the local + Hessian uses (``X_q`` when input quantization is enabled, ``X`` otherwise) with weight + fake-quant disabled, and normalized by the sample count. + + Like :func:`local_hessian_calibrate`, TensorQuantizer weights are calibrated — with the + weighted metric where a weight pairs with its input activations (dense linears and HF + fused-MoE experts), plain MSE otherwise. Other quantizer types (e.g. SequentialQuantizer) + are unsupported and left at their max-calibrated scale. + + Args: + model: Model to be calibrated. + forward_loop: A callable which takes the model as argument and + forwards calibration data through the model. Required for this algorithm. + distributed_sync: Whether to sync amax across distributed processes. + step_size: Step size for amax search (default: 0.1). + start_multiplier: Starting multiplier for amax search (default: 0.25). + stop_multiplier: Ending multiplier for amax search (default: 4.0). + fp8_scale_sweep: If True, sweep over all 128 possible FP8 E4M3 scale values + for NVFP4 per-block quantization (default: True). + block_size: Block size for the per-channel importance blocking (default: 16). + debug: If True, retain the per-quantizer importance accumulators on the model + (``model._local_hessian_accumulators``) for inspection. + + See :class:`WmseCalibConfig ` + for details on the configuration options. + """ + _hessian_weighted_calibrate( + model, + forward_loop, + distributed_sync, + step_size, + start_multiplier, + stop_multiplier, + fp8_scale_sweep, + block_size, + activation_error_coupling=False, + debug=debug, + shared_states=shared_states, + algorithm="wmse", + diagonal=True, + ) def enable_stats_collection(model: nn.Module): @@ -2273,6 +2403,7 @@ def _run_scale_calibration(model, forward_loop, scale_algorithm): calib_funcs = { "mse": mse_calibrate, "local_hessian": local_hessian_calibrate, + "wmse": wmse_calibrate, "max": max_calibrate, } calib_funcs[method](model, forward_loop=forward_loop, **algo_kwargs) @@ -2297,7 +2428,7 @@ def lsq( model: Quantized model. forward_loop: Calibration data forward loop. scale_algorithm: Calibration algorithm config to run first. - Dict with 'method' key: 'mse', 'local_hessian', or 'max'. + Dict with 'method' key: 'mse', 'local_hessian', 'wmse', or 'max'. Defaults to {'method': 'mse'} if None. learnable_amax: Which amax params are learnable: 'pre', 'post', ['pre', 'post'], or []. diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 1adda0511a6..4fc8a650976 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -259,6 +259,7 @@ def forward_loop(model) -> None: "max", "mse", "local_hessian", + "wmse", "smoothquant", "awq_lite", "awq_full", diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_wmse.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_wmse.yaml new file mode 100644 index 00000000000..a13e0b3cfac --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_wmse.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for NVFP4 W4A4 with static weight scales from weighted-MSE calibration. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + nvfp4_static: configs/numerics/nvfp4_static + +algorithm: + method: wmse + fp8_scale_sweep: true + layerwise: + enable: true +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*input_quantizer' + cfg: + $import: nvfp4 + - $import: default_disabled_quantizers diff --git a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py index 3b5b71845d8..bb68278f036 100644 --- a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py +++ b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py @@ -382,10 +382,10 @@ def forward_loop(m): def _build_hessian_accumulator( - cout, cin, quantizer_output, block_size=BLOCK_SIZE, unquantized_input=None + cout, cin, quantizer_output, block_size=BLOCK_SIZE, unquantized_input=None, diagonal=False ): """Real ``_LocalHessianAccumulator`` so the test exercises the production metric.""" - acc = _LocalHessianAccumulator(cout, cin, block_size) + acc = _LocalHessianAccumulator(cout, cin, block_size, diagonal=diagonal) acc.accumulate(quantizer_output, unquantized_input) return acc @@ -420,14 +420,17 @@ def _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc): def _total_hessian_loss(x_blocks, per_block_amax, global_amax, hessian): - """Total Hessian-weighted quantization error ``Σ dwᵀ H dw`` under the production - (CUDA ``static_blockwise_fp4_fake_quant``) rounding used at deployment — the objective - the sweep minimizes, summed over all blocks.""" + """Total weighted quantization error under the production (CUDA + ``static_blockwise_fp4_fake_quant``) rounding used at deployment — the objective the + sweep minimizes, summed over all blocks. ``Σ dwᵀ H dw`` for a full per-cin-block + Hessian, ``Σ Imp·dw²`` for the rank-2 diagonal (wmse).""" n_blocks = x_blocks.shape[0] n_cin = hessian.shape[0] h_per_block = hessian[torch.arange(n_blocks, device=x_blocks.device) % n_cin] xq = static_blockwise_fp4_fake_quant(x_blocks.float(), per_block_amax, global_amax) dw = x_blocks.float() - xq + if hessian.dim() == 2: + return ((dw * dw) * h_per_block).sum() return (torch.einsum("nij,nj->ni", h_per_block, dw) * dw).sum() @@ -488,6 +491,49 @@ def test_hessian_parity_random_weights(seed, cout, cin, dtype): ) +@requires_triton +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize(("cout", "cin"), [(8, 64), (1, 256), (256, 2048)]) +def test_wmse_diagonal_parity_random_weights(cout, cin, dtype): + """The diagonal (``wmse``) Triton sweep must match its reference 126-step sweep. + + Also pins the design invariant that a diagonal ``Imp`` selects exactly what a full + Hessian equal to ``diag(Imp)`` selects. + """ + torch.manual_seed(0) + device = "cuda" + weight = torch.randn(cout, cin, device=device, dtype=dtype) + activations = torch.randn(512, cin, device=device, dtype=torch.float32) + acc = _build_hessian_accumulator(cout, cin, activations, diagonal=True) + assert acc.normalized_hessian().shape == (cin // BLOCK_SIZE, BLOCK_SIZE) + + x_blocks = weight.reshape(-1, BLOCK_SIZE) + per_block_amax = x_blocks.float().abs().amax(dim=-1) + global_amax = per_block_amax.max() + + importance = acc.normalized_hessian() + ref = _run_hessian_reference(x_blocks, per_block_amax, global_amax, acc) + tri = _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc) + + assert ref.shape == tri.shape + n_blocks = ref.numel() + n_diff = int((ref != tri).sum()) + # The kernel and the reference einsum reduce the same 16 fp32 products in a different + # order, so a block whose two best candidates are exactly tied may flip. Cap that as the + # full-Hessian test does and require the achieved objective to be unchanged. + assert n_diff / n_blocks < 1e-3, f"{n_diff}/{n_blocks} blocks differ (>0.1%)" + loss_ref = _total_hessian_loss(x_blocks, ref, global_amax, importance) + loss_tri = _total_hessian_loss(x_blocks, tri, global_amax, importance) + rel_gap = ((loss_tri - loss_ref) / loss_ref.abs().clamp_min(1e-12)).abs().item() + assert rel_gap < 1e-6, f"aggregate wmse-loss gap {rel_gap:.3e} too large (dtype={dtype})" + + # diag(Imp) through the full-Hessian kernel must agree with the diagonal kernel. + embedded = _build_hessian_accumulator(cout, cin, activations) + embedded.hessian_per_block = torch.diag_embed(importance) + embedded.num_samples = 1 + assert torch.equal(tri, _run_hessian_triton(x_blocks, per_block_amax, global_amax, embedded)) + + @requires_triton def test_hessian_sweep_input_validation(): """``nvfp4_fp8_scale_sweep_hessian`` should reject malformed inputs cleanly.""" @@ -501,9 +547,11 @@ def test_hessian_sweep_input_validation(): nvfp4_fp8_scale_sweep_hessian(x.cpu(), g.cpu(), h.cpu()) with pytest.raises(ValueError, match="block_size"): nvfp4_fp8_scale_sweep_hessian(x, g, h, block_size=0) - # Wrong Hessian block dims. + # Wrong Hessian block dims (full and diagonal forms). with pytest.raises(ValueError, match="hessian must have shape"): nvfp4_fp8_scale_sweep_hessian(x, g, torch.randn(4, 8, 8, device=device)) + with pytest.raises(ValueError, match="hessian must have shape"): + nvfp4_fp8_scale_sweep_hessian(x, g, torch.randn(4, 8, device=device)) with pytest.raises(ValueError, match="coupling_bias must have"): nvfp4_fp8_scale_sweep_hessian(x, g, h, coupling_bias=torch.randn(3, device=device)) with pytest.raises(ValueError, match="coupling_bias must be a CUDA"): diff --git a/tests/unit/torch/quantization/test_wmse.py b/tests/unit/torch/quantization/test_wmse.py new file mode 100644 index 00000000000..62887b39f83 --- /dev/null +++ b/tests/unit/torch/quantization/test_wmse.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for weighted-MSE (``wmse``) calibration (CPU).""" + +import pytest +import torch +import torch.nn as nn +from _test_utils.torch.quantization.models import SimpleConv, SimpleLinear + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.config import WmseCalibConfig +from modelopt.torch.quantization.model_calib import ( + _LocalHessianAccumulator, + local_hessian_calibrate, + mse_calibrate, + wmse_calibrate, +) +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer + +# Weight-only INT8 per-channel; calibration is re-run explicitly per test. +INT8_WEIGHT_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", +} + +INT8_W8A8_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + ], + "algorithm": "max", +} + + +def _weight_amaxes(model): + return { + n: m.amax + for n, m in model.named_modules() + if isinstance(m, TensorQuantizer) and m.is_enabled and m.amax is not None + } + + +def _make_forward_loop(seed=0): + def forward_loop(model): + torch.manual_seed(seed) + for _ in range(3): + x = torch.randn(8, 16) + x[:, 0] *= 40.0 # skew so the importance is non-trivial vs plain weight MSE + model(x) + + return forward_loop + + +class TestWmseAccumulator: + def test_accumulate_shape_samples_fp32_buffer(self): + torch.manual_seed(0) + acc = _LocalHessianAccumulator(8, 32, 16, diagonal=True) + assert acc.is_enabled + acc.accumulate(torch.randn(10, 32, dtype=torch.bfloat16)) + assert acc.hessian_per_block.shape == (2, 16) # [n_blocks, block_size], not [n, bs, bs] + assert acc.hessian_per_block.dtype == torch.float32 # fp32 despite bf16 input + acc.accumulate(torch.randn(5, 32)) + assert acc.num_samples == 15 + assert acc.build_error_func() is not None + assert acc.hessian_per_block is None # raw buffer freed + + def test_importance_is_sum_of_squared_activations(self): + torch.manual_seed(1) + cin, bs = 32, 16 + x = torch.randn(7, cin) + acc = _LocalHessianAccumulator(4, cin, bs, diagonal=True) + acc.accumulate(x) + expected = x.square().sum(0).reshape(cin // bs, bs) # Imp_i = ‖X[:, i]‖² + assert torch.allclose(acc.hessian_per_block, expected, atol=1e-5) + assert torch.allclose(acc.normalized_hessian(), expected / 7, atol=1e-6) + + def test_importance_is_the_local_hessian_diagonal(self): + torch.manual_seed(2) + cin, bs = 32, 16 + x = torch.randn(9, cin) + full = _LocalHessianAccumulator(4, cin, bs) + diag = _LocalHessianAccumulator(4, cin, bs, diagonal=True) + full.accumulate(x) + diag.accumulate(x) + assert torch.allclose( + diag.hessian_per_block, torch.diagonal(full.hessian_per_block, dim1=-2, dim2=-1) + ) + + def test_error_func_matches_explicit_weighted_squared_error(self): + torch.manual_seed(3) + cout, cin, bs = 4, 32, 16 + n_blocks = cin // bs + acc = _LocalHessianAccumulator(cout, cin, bs, diagonal=True) + x = torch.randn(7, cin) + acc.accumulate(x) + error_func = acc.build_error_func() + + importance = x.square().sum(0).reshape(n_blocks, bs) / acc.num_samples + w = torch.randn(cout * n_blocks, bs) + wq = w + 0.05 * torch.randn_like(w) + err = error_func(w, wq).view(-1, bs) + + assert err.shape == (cout * n_blocks, bs) + assert torch.allclose(err, err[:, :1].expand(-1, bs)) # per-block scalar broadcast + dw = (w - wq).view(cout, n_blocks, bs) + expected = torch.einsum("cnb,nb->cn", dw * dw, importance).reshape(-1) + assert torch.allclose(err[:, 0], expected, atol=1e-5) + + def test_error_func_matches_local_hessian_with_diagonal_hessian(self): + """Pins eq. 13: WMSE(Imp) == local_hessian(diag(Imp)) block-for-block.""" + torch.manual_seed(4) + cout, cin, bs = 5, 32, 16 + n_blocks = cin // bs + x = torch.randn(11, cin) + + diag_acc = _LocalHessianAccumulator(cout, cin, bs, diagonal=True) + diag_acc.accumulate(x) + importance = diag_acc.normalized_hessian() + wmse_error = diag_acc.build_error_func() + + # A full-Hessian accumulator whose Hessian is exactly diag(Imp). + full_acc = _LocalHessianAccumulator(cout, cin, bs) + full_acc.hessian_per_block = torch.diag_embed(importance) + full_acc.num_samples = 1 + lh_error = full_acc.build_error_func() + + w = torch.randn(cout * n_blocks, bs) + wq = w + 0.05 * torch.randn_like(w) + assert torch.allclose(wmse_error(w, wq), lh_error(w, wq), atol=1e-5) + + def test_returns_none_when_disabled_or_no_samples(self): + not_divisible = _LocalHessianAccumulator(8, 30, 16, diagonal=True) + assert not not_divisible.is_enabled + not_divisible.accumulate(torch.randn(4, 30)) # no-op + assert not_divisible.build_error_func() is None + # no samples + assert _LocalHessianAccumulator(8, 32, 16, diagonal=True).build_error_func() is None + + def test_wmse_never_allocates_coupling(self): + """``activation_error_coupling`` is local-Hessian-only; wmse must not carry it.""" + acc = _LocalHessianAccumulator(2, 4, 2, diagonal=True) + acc.accumulate(torch.randn(3, 4)) + assert acc.coupling_per_block is None + assert acc.normalized_coupling() is None + + +def test_wmse_config_and_preset(): + config = WmseCalibConfig() + dumped = config.model_dump() + assert dumped["method"] == "wmse" + assert dumped["fp8_scale_sweep"] is True + assert dumped["block_size"] == 16 + assert "activation_error_coupling" not in dumped # local-Hessian-only feature + with pytest.raises(ValueError): + WmseCalibConfig(activation_error_coupling=True) + + preset = mtq.NVFP4_W4A4_WEIGHT_WMSE_CFG + assert preset["algorithm"]["method"] == "wmse" + assert preset["algorithm"]["fp8_scale_sweep"] is True + assert preset["algorithm"]["layerwise"]["enable"] is True + # Same numerics as the local-Hessian preset; only the scale search differs. + assert preset["quant_cfg"] == mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG["quant_cfg"] + + +class TestWmseCalibrateDense: + def test_refines_amax_beyond_max_and_plain_mse(self): + forward_loop = _make_forward_loop() + torch.manual_seed(0) + model_wmse = SimpleLinear() + mtq.quantize(model_wmse, INT8_WEIGHT_CFG, forward_loop=forward_loop) + max_amax = {n: a.clone() for n, a in _weight_amaxes(model_wmse).items()} + wmse_calibrate(model_wmse, forward_loop, fp8_scale_sweep=False, debug=True) + + torch.manual_seed(0) + model_mse = SimpleLinear() + mtq.quantize(model_mse, INT8_WEIGHT_CFG, forward_loop=forward_loop) + mse_calibrate(model_mse, forward_loop, fp8_scale_sweep=False) + + accs = model_wmse._local_hessian_accumulators + assert accs and all(a.num_samples > 0 for a in accs.values()) + assert all(a.diagonal and a.normalized_hessian().dim() == 2 for a in accs.values()) + wmse, mse = _weight_amaxes(model_wmse), _weight_amaxes(model_mse) + assert all(torch.isfinite(a).all() and (a > 0).all() for a in wmse.values()) + assert any(not torch.allclose(wmse[n], max_amax[n]) for n in wmse) # refined past max-cal + assert any(not torch.allclose(wmse[n], mse[n]) for n in wmse) # weighting changed choice + + def test_matches_local_hessian_when_hessian_is_diagonal(self): + """End-to-end equivalence: a diagonal input covariance makes the two agree exactly. + + Single linear so the calibrated layer's own input is the crafted one; one-hot rows + scaled per channel give ``XᵀX = diag(scale²)`` exactly, so the local Hessian carries + no information the diagonal importance does not and both must pick the same amax. + """ + + class _OneLinear(nn.Module): + def __init__(self): + super().__init__() + self.fc = nn.Linear(16, 8) + + def forward(self, x): + return self.fc(x) + + scale = torch.logspace(-1, 1, 16) + x = torch.eye(16) * scale + forward_loop = lambda m: m(x) # noqa: E731 + + amaxes = {} + for name, calibrate in (("wmse", wmse_calibrate), ("lh", local_hessian_calibrate)): + torch.manual_seed(0) + model = _OneLinear() + mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=forward_loop) + calibrate(model, forward_loop, fp8_scale_sweep=False) + amaxes[name] = _weight_amaxes(model) + + assert amaxes["wmse"] and amaxes["wmse"].keys() == amaxes["lh"].keys() + assert all(torch.allclose(amaxes["wmse"][n], amaxes["lh"][n]) for n in amaxes["wmse"]) + + def test_warns_with_module_name_when_cin_not_divisible(self): + class _OddModel(nn.Module): + def __init__(self): + super().__init__() + self.odd = nn.Linear(24, 32) # 24 not divisible by block_size 16 + + def forward(self, x): + return self.odd(x) + + torch.manual_seed(0) + model = _OddModel() + forward_loop = lambda m: m(torch.randn(4, 24)) # noqa: E731 + mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=forward_loop) + with pytest.warns(UserWarning, match=r"wmse: odd input features \(24\) not divisible"): + wmse_calibrate(model, forward_loop, fp8_scale_sweep=False) + + def test_no_forward_loop_is_skipped(self): + torch.manual_seed(0) + model = SimpleLinear() + mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=_make_forward_loop()) + before = {n: a.clone() for n, a in _weight_amaxes(model).items()} + with pytest.warns(UserWarning, match="forward_loop must be provided for wmse"): + wmse_calibrate(model, forward_loop=None) + assert all(torch.equal(before[n], a) for n, a in _weight_amaxes(model).items()) + + @pytest.mark.parametrize("quant_cfg", [INT8_WEIGHT_CFG, INT8_W8A8_CFG]) + def test_importance_uses_input_quantizer_output(self, quant_cfg): + torch.manual_seed(0) + model = SimpleLinear() + x = torch.randn(8, 16) + x[:, 0] *= 40.0 + forward_loop = lambda m: m(x) # noqa: E731 + mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + wmse_calibrate(model, forward_loop, fp8_scale_sweep=False, debug=True) + + linear = model.net[0] + acc = model._local_hessian_accumulators[id(linear.weight_quantizer)] + quantizer_output = linear.input_quantizer(x).float() + expected = quantizer_output.square().sum(0).reshape(1, 16) + assert torch.equal(acc.hessian_per_block, expected) + assert acc.coupling_per_block is None + + +class TestWmseFallbacks: + """Weights wmse can't pair with an input fall back to plain MSE (no importance).""" + + def test_conv_weight_falls_back_without_crash(self): + torch.manual_seed(0) + model = SimpleConv() # 4-D conv weights — no single 2-D weight to pair + forward_loop = lambda m: m(SimpleConv.get_input()) # noqa: E731 + mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=forward_loop) + wmse_calibrate(model, forward_loop, fp8_scale_sweep=False, debug=True) + conv = model.net[0] + assert id(conv.weight_quantizer) not in model._local_hessian_accumulators + assert conv.weight_quantizer.amax is not None # still calibrated via plain MSE + + def test_sequential_quantizer_weight_falls_back_without_crash(self): + torch.manual_seed(0) + model = SimpleLinear() + mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=_make_forward_loop()) + linear = model.net[0] + linear.weight_quantizer = SequentialQuantizer(TensorQuantizer(), TensorQuantizer()) + wmse_calibrate(model, _make_forward_loop(), fp8_scale_sweep=False, debug=True) + assert id(linear.weight_quantizer) not in model._local_hessian_accumulators