diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 012c6512b29..8e6b14e2468 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -24,6 +24,7 @@ Changelog - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv: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: diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 3a7380a9f9e..594ec30d792 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -162,7 +162,7 @@ export HF_PATH= --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`.* diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py index a56588b21d8..fd6cd0d389b 100644 --- a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py +++ b/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 diff --git a/modelopt/torch/quantization/calib/mse.py b/modelopt/torch/quantization/calib/mse.py index 7d6583aacd2..88dd64b7a36 100644 --- a/modelopt/torch/quantization/calib/mse.py +++ b/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 diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 6c07afab90c..c8eb109ca0c 100644 --- a/modelopt/torch/quantization/config.py +++ b/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", } diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 031e766c687..21fd9429d14 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -21,7 +21,7 @@ import warnings from collections.abc import Callable, Mapping, Sequence from functools import partial -from typing import TypeAlias +from typing import NamedTuple, TypeAlias import torch import torch.distributed as dist @@ -523,12 +523,13 @@ def _make_weight_mse_calibrator( stop_multiplier: float, fp8_scale_sweep: bool, error_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - hessian: torch.Tensor | None = None, + metric: "_LocalHessianMetric | None" = None, ) -> _Calibrator | None: """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). - ``hessian`` (the same per-cin-block metric as a raw tensor) enables NVFP4's Hessian-weighted + ``metric`` carries the per-cin-block Hessian and optional activation error coupling matrix, + enabling NVFP4's Hessian-weighted Triton fast path; ``error_func`` then serves only as the reference fallback. """ if ( @@ -565,7 +566,8 @@ def _make_weight_mse_calibrator( global_amax=weight_quantizer.global_amax, quant_func=quant_func, error_func=error_func, - hessian=hessian, + hessian=metric.hessian if metric is not None else None, + coupling=metric.coupling if metric is not None else None, ) # fp8_scale_sweep applies only to registered backends and static NVFP4; skip others. return None @@ -638,14 +640,13 @@ def _mse_calibrate_weights( stop_multiplier: float, fp8_scale_sweep: bool, error_func_for: Callable[[TensorQuantizer], Callable | None] | None = None, - hessian_for: Callable[[TensorQuantizer], torch.Tensor | None] | None = None, + metric_for: Callable[[TensorQuantizer], "_LocalHessianMetric | None"] | None = None, ): """Run MSE weight calibration over all eligible quantizers (shared by mse / local-Hessian). ``error_func_for`` maps a weight quantizer to an optional per-weight error function - (local-Hessian's Hessian metric); ``None`` means plain squared error. ``hessian_for`` - maps a weight quantizer to the same metric as a raw per-cin-block Hessian tensor, - enabling the Hessian-weighted Triton fast path. + (local-Hessian's metric); ``None`` means plain squared error. ``metric_for`` maps a weight + quantizer to the Hessian and optional activation error coupling matrix used by the Triton path. """ seen_modules: set[int] = set() pbar = tqdm(desc="MSE weight calibration") @@ -656,7 +657,7 @@ def _mse_calibrate_weights( with enable_weight_access_and_writeback(parent_module, model, name_to_module): for weight, weight_quantizer in parent_module.iter_weights_for_calibration(): error_func = error_func_for(weight_quantizer) if error_func_for else None - hessian = hessian_for(weight_quantizer) if hessian_for else None + metric = metric_for(weight_quantizer) if metric_for else None cal = _make_weight_mse_calibrator( weight_quantizer, step_size, @@ -664,7 +665,7 @@ def _mse_calibrate_weights( stop_multiplier, fp8_scale_sweep, error_func=error_func, - hessian=hessian, + metric=metric, ) if cal is None: continue @@ -679,11 +680,44 @@ def _mse_calibrate_weights( pbar.close() +class _LocalHessianMetric(NamedTuple): + hessian: torch.Tensor + coupling: torch.Tensor | None = None + + +def _activation_error_coupling_bias( + original_weight: torch.Tensor, + coupling: torch.Tensor, + cout: int, + block_size: int, +) -> torch.Tensor: + """Return ``P @ W0`` for the activation error coupling term. + + With ``Δ = quantized - original`` and asymmetric + ``P[i,d] = X_q[:,i]ᵀ(X_q-X)[:,d] / B``, the linear objective is + ``2 Σ_i Σ_d ΔW[i] P[i,d] W0[d]``. The returned ``[cout, n_blocks, block_size]`` + tensor is computed in fp32 and may produce a negative total objective after the + scale-independent activation-error constant is dropped. + """ + weight = original_weight.reshape(cout, -1, block_size).to(torch.float32) + coupling = coupling.to(device=weight.device, dtype=torch.float32) + if weight.is_cuda: + allow_tf32 = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = False + return torch.einsum("nbd,cnd->cnb", coupling, weight) + finally: + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + return torch.einsum("nbd,cnd->cnb", coupling, weight) + + class _LocalHessianAccumulator: - """Per-block local Hessian ``H = ΣXᵀX`` for one weight quantizer. + """Per-block local Hessian from input-quantizer outputs for one weight quantizer. Partitioned over ``cin`` into ``cin // block_size`` blocks to match the NVFP4 per-block - scale; the buffer is allocated lazily so never-routed experts cost nothing. + 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. """ def __init__(self, cout: int, cin: int, block_size: int): @@ -694,23 +728,39 @@ def __init__(self, cout: int, cin: int, block_size: int): # Not block-divisible -> no Hessian (falls back to plain MSE). self.is_enabled = cin % block_size == 0 self.hessian_per_block: torch.Tensor | None = None + self.coupling_per_block: torch.Tensor | None = None self._normalized_hessian: torch.Tensor | None = None + self._normalized_coupling: torch.Tensor | None = None self.num_samples = 0 @torch.no_grad() - def accumulate(self, input_tensor: torch.Tensor) -> None: - """Accumulate ``XᵀX`` per block from an activation of shape ``(..., cin)``.""" + def accumulate( + self, quantizer_output: torch.Tensor, unquantized_input: torch.Tensor | None = None + ) -> None: + """Accumulate the Hessian and optional coupling to ``unquantized_input``.""" if not self.is_enabled: return # fp32 GEMM avoids bf16/fp16 precision loss; (cin, tokens) -> (n_blocks, bs, tokens). - x = input_tensor.reshape(-1, self.cin).to(torch.float32).T + x = quantizer_output.reshape(-1, self.cin).to(torch.float32).T + dx = ( + x - unquantized_input.reshape(-1, self.cin).to(torch.float32).T + if unquantized_input is not None and unquantized_input is not quantizer_output + else None + ) x = x.reshape(self.num_blocks_per_cin, self.block_size, -1) hessian_batch = x @ x.transpose(-1, -2) if self.hessian_per_block is None: self.hessian_per_block = hessian_batch else: self.hessian_per_block += hessian_batch - self.num_samples += input_tensor.numel() // self.cin + if dx is not None: + dx = dx.reshape(self.num_blocks_per_cin, self.block_size, -1) + coupling_batch = x @ dx.transpose(-1, -2) + if self.coupling_per_block is None: + self.coupling_per_block = coupling_batch + else: + self.coupling_per_block += coupling_batch + 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). @@ -726,6 +776,16 @@ def normalized_hessian(self) -> torch.Tensor | None: self._normalized_hessian = self.hessian_per_block / self.num_samples return self._normalized_hessian + def normalized_coupling(self) -> torch.Tensor | None: + """Normalized activation error coupling matrix, or ``None`` when not captured.""" + if ( + self._normalized_coupling is None + and self.coupling_per_block is not None + and self.num_samples + ): + self._normalized_coupling = self.coupling_per_block / self.num_samples + return self._normalized_coupling + def build_error_func( self, keep_buffer: bool = False ) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None: @@ -734,18 +794,24 @@ def build_error_func( Frees the raw Hessian buffer unless ``keep_buffer`` (kept for debug inspection). """ hessian = self.normalized_hessian() + coupling = self.normalized_coupling() if hessian is None: return None cout = self.cout bs = self.block_size if not keep_buffer: self.hessian_per_block = None + self.coupling_per_block = None 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 = (x - xq).view(cout, -1, bs) - block_loss = torch.einsum("cnb,nbd,cnd->cn", dw, hessian, dw).reshape(-1) + dw = (xq - x).view(cout, -1, bs) + block_loss = 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) + block_loss = block_loss.reshape(-1) return block_loss.unsqueeze(-1).expand(-1, bs).reshape(original_shape) return local_hessian_error @@ -786,7 +852,24 @@ def _is_quant_fused_experts(module: nn.Module) -> bool: ) -def _register_local_hessian_input_hooks(model, name_to_module, capture, block_size, warned): +def _activation_error_coupling_fallback_reason(input_quantizer) -> str | None: + if not isinstance(input_quantizer, TensorQuantizer) or not input_quantizer.is_enabled: + return None + if input_quantizer.pre_quant_scale is not None: + return "input quantizer uses pre_quant_scale" + if input_quantizer.rotate_is_enabled or input_quantizer.rotate_back_is_enabled: + return "input quantizer uses rotation" + return "" + + +def _register_local_hessian_input_hooks( + model, + name_to_module, + capture, + block_size, + warned, + activation_error_coupling: bool = False, +): """Register forward hooks feeding each weight's input activations to ``capture``. Local-Hessian-specific (kept here rather than as a general ``QuantModule`` API): dense @@ -796,17 +879,35 @@ def _register_local_hessian_input_hooks(model, name_to_module, capture, block_si """ handles: list = [] - def _make_expert_hook(expert_module, weight_name, quantizers, enabled): - def _expert_hook(_input_quantizer, args): + def _make_expert_hook(expert_module, weight_name, quantizers, enabled, capture_coupling): + def _expert_hook(_input_quantizer, args, output): if not args: return idx = expert_module._current_expert_idx if idx in enabled: # Read the weight fresh (valid under accelerate/FSDP re-materialization). - capture(quantizers[idx], getattr(expert_module, weight_name)[idx], args[0]) + capture( + quantizers[idx], + getattr(expert_module, weight_name)[idx], + output, + args[0] if capture_coupling else None, + ) return _expert_hook + def _capture_coupling(input_quantizer, layer_name): + if not activation_error_coupling: + return False + reason = _activation_error_coupling_fallback_reason(input_quantizer) + if reason == "": + return True + if reason is not None: + warn_rank_0( + "activation error coupling falls back to the Hessian-only objective for " + f"{layer_name}: {reason}." + ) + return False + for name, module in name_to_module.items(): if is_quantized_linear(module) and isinstance(module.weight_quantizer, TensorQuantizer): with enable_weight_access_and_writeback(module, model, name_to_module): @@ -819,11 +920,23 @@ def _expert_hook(_input_quantizer, args): name, weight, module.weight_quantizer, block_size, warned ) - def _dense_hook(linear, args): + input_quantizer = getattr(module, "input_quantizer", None) + if not isinstance(input_quantizer, nn.Module): + continue + capture_coupling = _capture_coupling(input_quantizer, name) + + def _dense_hook( + _input_quantizer, args, output, linear=module, coupling=capture_coupling + ): if args: - capture(linear.weight_quantizer, linear.weight, args[0]) + capture( + linear.weight_quantizer, + linear.weight, + output, + args[0] if coupling else None, + ) - handles.append(module.register_forward_pre_hook(_dense_hook)) + handles.append(input_quantizer.register_forward_hook(_dense_hook)) elif _is_quant_fused_experts(module): with enable_weight_access_and_writeback(module, model, name_to_module): first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") @@ -846,9 +959,17 @@ def _dense_hook(linear, args): # Snapshot which experts are enabled now, before the caching forward silences # all weight quantizers — so we don't capture (and discard) disabled experts. enabled = {i for i, q in enumerate(quantizers) if q.is_enabled} + layer_name = f"{name}.{weight_name}" + capture_coupling = _capture_coupling(input_quantizer, layer_name) handles.append( - input_quantizer.register_forward_pre_hook( - _make_expert_hook(module, weight_name, quantizers, enabled) + input_quantizer.register_forward_hook( + _make_expert_hook( + module, + weight_name, + quantizers, + enabled, + capture_coupling, + ) ) ) return handles @@ -864,20 +985,29 @@ def local_hessian_calibrate( 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 ``(W - Wq)ᵀ H (W - Wq)`` with per-block Hessian ``H = ΣXᵀX`` (approximating the - output error ``||WX - WqX||²``), built from a forward with weight fake-quant disabled - (input quantizers untouched) and fed to :func:`mse_calibrate`'s weight search via ``error_func``. + 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 @@ -889,6 +1019,8 @@ def local_hessian_calibrate( 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. @@ -908,19 +1040,33 @@ def local_hessian_calibrate( # Hessians keyed by id(weight_quantizer); modules pair weights<->activations via the hook. accumulators: dict[int, _LocalHessianAccumulator] = {} - def capture(weight_quantizer, weight, input_tensor): - input_local = input_tensor.to_local() if hasattr(input_tensor, "to_local") else input_tensor + def capture(weight_quantizer, weight, quantizer_output, unquantized_input=None): + quantizer_output_local = ( + quantizer_output.to_local() + if hasattr(quantizer_output, "to_local") + else quantizer_output + ) + unquantized_input_local = ( + unquantized_input.to_local() + if unquantized_input is not None and hasattr(unquantized_input, "to_local") + else unquantized_input + ) acc = accumulators.get(id(weight_quantizer)) if acc is None: acc = _LocalHessianAccumulator(weight.shape[0], weight.shape[1], block_size) accumulators[id(weight_quantizer)] = acc - acc.accumulate(input_local) + acc.accumulate(quantizer_output_local, unquantized_input_local) - # Phase 2: capture each weight's input activations during a forward with weight fake-quant - # disabled (so H = ΣXᵀX reflects full-precision weights); input quantizers are left as-is. + # Phase 2: capture each input quantizer's output during a forward with weight fake-quant + # disabled; enabled input quantizers contribute X_q to H and disabled ones contribute X. warned: set = set() handles = _register_local_hessian_input_hooks( - model, name_to_module, capture, block_size, warned + model, + name_to_module, + capture, + block_size, + warned, + activation_error_coupling=activation_error_coupling, ) print_rank_0("local_hessian: Caching activations and computing local Hessian...") try: @@ -946,7 +1092,11 @@ def capture(weight_quantizer, weight, input_tensor): error_funcs = { qid: acc.build_error_func(keep_buffer=debug) for qid, acc in accumulators.items() } - hessians = {qid: acc.normalized_hessian() for qid, acc in accumulators.items()} + metrics = { + qid: _LocalHessianMetric(acc.normalized_hessian(), acc.normalized_coupling()) + 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...") _mse_calibrate_weights( model, @@ -956,18 +1106,19 @@ def capture(weight_quantizer, weight, input_tensor): stop_multiplier=stop_multiplier, fp8_scale_sweep=fp8_scale_sweep, error_func_for=lambda q: error_funcs.get(id(q)), - hessian_for=lambda q: hessians.get(id(q)), + metric_for=lambda q: metrics.get(id(q)), ) # Release the per-block Hessians (held by the error_func closures, calibrators, and the # accumulators' cache) before empty_cache so export starts defragmented; keep only for debug. error_funcs.clear() - hessians.clear() + metrics.clear() for module in name_to_module.values(): if isinstance(module, TensorQuantizer) and isinstance(module._calibrator, MseCalibrator): module._calibrator._error_func = None if isinstance(module._calibrator, NVFP4MSECalibrator): module._calibrator._hessian = None + module._calibrator._coupling = None if debug: model._local_hessian_accumulators = accumulators else: diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian.yaml index ac6a3094b7c..413dc8d5ad3 100644 --- a/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian.yaml +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian.yaml @@ -25,6 +25,8 @@ imports: algorithm: method: local_hessian fp8_scale_sweep: true + layerwise: + enable: true quant_cfg: - $import: base_disable_all - quantizer_name: '*weight_quantizer' diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_error_coupling.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_error_coupling.yaml new file mode 100644 index 00000000000..78fa97926e8 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_error_coupling.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# QuantizeConfig preset for NVFP4 W4A4 local-Hessian calibration with activation error coupling. + +# 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: local_hessian + fp8_scale_sweep: true + activation_error_coupling: 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 2c309f5a532..3b5b71845d8 100644 --- a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py +++ b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py @@ -381,10 +381,12 @@ def forward_loop(m): # -------------------------------------------------------------------------------------- -def _build_hessian_accumulator(cout, cin, hessian_input, block_size=BLOCK_SIZE): +def _build_hessian_accumulator( + cout, cin, quantizer_output, block_size=BLOCK_SIZE, unquantized_input=None +): """Real ``_LocalHessianAccumulator`` so the test exercises the production metric.""" acc = _LocalHessianAccumulator(cout, cin, block_size) - acc.accumulate(hessian_input) + acc.accumulate(quantizer_output, unquantized_input) return acc @@ -411,6 +413,7 @@ def _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc): global_amax=global_amax, quant_func=_reference_quant_func(global_amax), hessian=acc.normalized_hessian(), + coupling=acc.normalized_coupling(), ) cal.collect(x_blocks) return cal.compute_amax() @@ -501,6 +504,30 @@ def test_hessian_sweep_input_validation(): # Wrong Hessian block dims. 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="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"): + nvfp4_fp8_scale_sweep_hessian(x, g, h, coupling_bias=torch.randn(x.numel())) + + +@requires_triton +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_activation_error_coupling_hessian_parity(dtype): + """Activation error coupling selects equivalent scales in reference and Triton paths.""" + torch.manual_seed(17) + cout, cin = 32, 128 + weight = torch.randn(cout, cin, device="cuda", dtype=dtype) + x = torch.randn(256, cin, device="cuda") + xq = (x / 0.125).round() * 0.125 + acc = _build_hessian_accumulator(cout, cin, xq, unquantized_input=x) + x_blocks = weight.reshape(-1, BLOCK_SIZE) + per_block_amax = x_blocks.float().abs().amax(dim=-1) + global_amax = per_block_amax.max() + + ref = _run_hessian_reference(x_blocks, per_block_amax, global_amax, acc) + tri = _run_hessian_triton(x_blocks, per_block_amax, global_amax, acc) + mismatch = (ref != tri).float().mean().item() + assert mismatch < 1e-3 @requires_triton diff --git a/tests/gpu/torch/quantization/test_quantize_cuda.py b/tests/gpu/torch/quantization/test_quantize_cuda.py index 8fb4c096030..77e88af3e02 100644 --- a/tests/gpu/torch/quantization/test_quantize_cuda.py +++ b/tests/gpu/torch/quantization/test_quantize_cuda.py @@ -97,6 +97,7 @@ mtq.NVFP4_AWQ_CLIP_CFG, mtq.NVFP4_AWQ_FULL_CFG, mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG, + mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_ACT_ERROR_COUPLING_CFG, mtq.MXFP8_DEFAULT_CFG, mtq.MXFP6_DEFAULT_CFG, mtq.MXFP4_DEFAULT_CFG, @@ -124,6 +125,7 @@ def test_quantize(model_cls, config): NVFP4_WEIGHT_ACT_MSE_CFG, NVFP4_WEIGHT_MSE_FP8_SWEEP_CFG, mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG, + mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_ACT_ERROR_COUPLING_CFG, ]: if get_cuda_ext_mx() is None: pytest.skip("cuda_ext_mx is not available") diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 9fa836bb620..7c15e285513 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -899,7 +899,8 @@ def forward_loop(m): self._cleanup_registry(expert_type) - def test_local_hessian_refines_per_expert_weights(self): + @pytest.mark.parametrize("activation_error_coupling", [False, True]) + def test_local_hessian_refines_per_expert_weights(self, activation_error_coupling): """local_hessian captures each expert's routed activations and refines its weight amax.""" model = _TinyMoEModel() expert_type = type(model.moe.experts) @@ -911,6 +912,10 @@ def test_local_hessian_refines_per_expert_weights(self): {"quantizer_name": "*", "enable": False}, {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": weight_quant}, {"quantizer_name": "*down_proj_weight_quantizer", "cfg": weight_quant}, + { + "quantizer_name": "*input_quantizer", + "cfg": {"num_bits": 8, "axis": None}, + }, ], "algorithm": "max", } @@ -936,7 +941,13 @@ def forward_loop(m): for i, q in enumerate(quantizers): expected_shape[id(q)] = (weight[i].shape[0], weight[i].shape[1]) - local_hessian_calibrate(model, forward_loop, fp8_scale_sweep=False, debug=True) + local_hessian_calibrate( + model, + forward_loop, + fp8_scale_sweep=False, + activation_error_coupling=activation_error_coupling, + debug=True, + ) # Each captured Hessian is keyed to a real per-expert quantizer with the matching weight # shape, spans multiple distinct experts, and the refinement moved at least one amax. @@ -944,6 +955,7 @@ def forward_loop(m): assert len(routed) >= 2, "expected multiple distinct experts to capture Hessians" for qid, acc in routed.items(): assert (acc.cout, acc.cin) == expected_shape[qid] + assert (acc.coupling_per_block is not None) is activation_error_coupling assert all(q.amax is not None and torch.isfinite(q.amax).all() for q in expert_quantizers) assert any( id(q) in max_amax and not torch.allclose(q.amax, max_amax[id(q)]) diff --git a/tests/unit/torch/quantization/test_local_hessian.py b/tests/unit/torch/quantization/test_local_hessian.py index d8b4c4eecec..b6f40f3a1fa 100644 --- a/tests/unit/torch/quantization/test_local_hessian.py +++ b/tests/unit/torch/quantization/test_local_hessian.py @@ -24,9 +24,11 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization import calib -from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.config import LocalHessianCalibConfig, QuantizerAttributeConfig from modelopt.torch.quantization.model_calib import ( _FP8_SWEEP_CALIBRATOR_REGISTRY, + _activation_error_coupling_bias, + _activation_error_coupling_fallback_reason, _LocalHessianAccumulator, _make_weight_mse_calibrator, _register_fp8_sweep_calibrator, @@ -49,6 +51,15 @@ "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 { @@ -110,6 +121,61 @@ def test_returns_none_when_disabled_or_no_samples(self): assert not_divisible.build_error_func() is None assert _LocalHessianAccumulator(8, 32, 16).build_error_func() is None # no samples + def test_activation_error_coupling_accumulation(self): + torch.manual_seed(2) + cout, cin, bs = 3, 8, 4 + x = torch.randn(6, cin, dtype=torch.bfloat16) + xq = (x.float() + 0.1 * torch.randn_like(x.float())).to(torch.bfloat16) + acc = _LocalHessianAccumulator(cout, cin, bs) + acc.accumulate(xq, x) + + xb = x.float().T.reshape(cin // bs, bs, -1) + xqb = xq.float().T.reshape(cin // bs, bs, -1) + assert acc.coupling_per_block.dtype == torch.float32 + assert torch.equal(acc.hessian_per_block, xqb @ xqb.transpose(-1, -2)) + assert torch.equal(acc.coupling_per_block, xqb @ (xqb - xb).transpose(-1, -2)) + + def test_activation_error_coupling_objective_matches_output_mse(self): + torch.manual_seed(3) + samples, cout, cin, bs = 7, 2, 2, 2 + xq = torch.randn(samples, cin) + x = xq - 0.2 * torch.randn_like(xq) + w0 = torch.randn(cout, cin) + acc = _LocalHessianAccumulator(cout, cin, bs) + acc.accumulate(xq, x) + hessian = acc.normalized_hessian()[0] + coupling = acc.normalized_coupling() + bias = _activation_error_coupling_bias(w0, coupling, cout, bs)[:, 0] + # The minimizer of the constant-dropped quadratic is deliberately negative. + dw = -torch.linalg.solve(hessian, bias.T).T + wq = w0 + dw + + objective = acc.build_error_func(keep_buffer=True)(w0, wq)[:, 0] + deployed = xq @ wq.T + reference = x @ w0.T + constant = xq @ w0.T - reference + expected = ((deployed - reference).square() - constant.square()).sum(0) / samples + assert torch.allclose(objective, expected, atol=2e-5, rtol=2e-5) + assert (objective < 0).all() + + def test_default_accumulation_does_not_allocate_coupling(self): + acc = _LocalHessianAccumulator(2, 4, 2) + x = torch.randn(3, 4) + acc.accumulate(x) + assert acc.coupling_per_block is None + assert acc.normalized_coupling() is None + + +def test_local_hessian_config_and_presets(): + config = LocalHessianCalibConfig(activation_error_coupling=True) + assert config.model_dump()["activation_error_coupling"] is True + presets = [ + mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG, + mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_ACT_ERROR_COUPLING_CFG, + ] + assert all(preset["algorithm"]["layerwise"]["enable"] is True for preset in presets) + assert presets[1]["algorithm"]["activation_error_coupling"] is True + class TestLocalHessianCalibrateDense: def test_refines_amax_beyond_max_and_plain_mse(self): @@ -157,6 +223,48 @@ def test_no_forward_loop_is_skipped(self): local_hessian_calibrate(model, forward_loop=None) assert all(torch.equal(before[n], a) for n, a in _weight_amaxes(model).items()) + def test_activation_error_coupling_captures_quantized_inputs(self): + torch.manual_seed(0) + model = SimpleLinear() + forward_loop = _make_forward_loop() + mtq.quantize(model, INT8_W8A8_CFG, forward_loop=forward_loop) + local_hessian_calibrate( + model, + forward_loop, + fp8_scale_sweep=False, + activation_error_coupling=True, + debug=True, + ) + routed = [a for a in model._local_hessian_accumulators.values() if a.num_samples] + assert routed + assert all(a.coupling_per_block is not None for a in routed) + assert any(torch.count_nonzero(a.coupling_per_block) for a in routed) + + @pytest.mark.parametrize("quant_cfg", [INT8_WEIGHT_CFG, INT8_W8A8_CFG]) + @pytest.mark.parametrize("activation_error_coupling", [False, True]) + def test_hessian_uses_input_quantizer_output(self, quant_cfg, activation_error_coupling): + 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) + local_hessian_calibrate( + model, + forward_loop, + fp8_scale_sweep=False, + activation_error_coupling=activation_error_coupling, + debug=True, + ) + + linear = model.net[0] + acc = model._local_hessian_accumulators[id(linear.weight_quantizer)] + quantizer_output = linear.input_quantizer(x).float().T.reshape(1, 16, -1) + expected_hessian = quantizer_output @ quantizer_output.transpose(-1, -2) + assert torch.equal(acc.hessian_per_block, expected_hessian) + expect_coupling = activation_error_coupling and linear.input_quantizer.is_enabled + assert (acc.coupling_per_block is not None) is expect_coupling + class TestLocalHessianFallbacks: """Weights local-Hessian can't pair with an input fall back to plain MSE (no Hessian).""" @@ -180,6 +288,25 @@ def test_sequential_quantizer_weight_falls_back_without_crash(self): local_hessian_calibrate(model, _make_forward_loop(), fp8_scale_sweep=False, debug=True) assert id(linear.weight_quantizer) not in model._local_hessian_accumulators + @pytest.mark.parametrize( + ("cfg", "reason"), + [ + ({"num_bits": 8}, "pre_quant_scale"), + ( + { + "num_bits": 8, + "rotate": {"enable": True, "mode": "rotate", "block_size": 8}, + }, + "rotation", + ), + ], + ) + def test_activation_error_coupling_fallback_reasons(self, cfg, reason): + q = TensorQuantizer(QuantizerAttributeConfig(**cfg)) + if reason == "pre_quant_scale": + q.pre_quant_scale = torch.ones(16) + assert reason in _activation_error_coupling_fallback_reason(q) + class TestBlockSizeMismatchWarning: def test_block_size_mismatch_warns_only_on_mismatch(self): diff --git a/tests/unit/torch/quantization/test_lsq.py b/tests/unit/torch/quantization/test_lsq.py index 50fd1ffc870..589adbf5c43 100644 --- a/tests/unit/torch/quantization/test_lsq.py +++ b/tests/unit/torch/quantization/test_lsq.py @@ -120,6 +120,12 @@ def test_unsupported_scale_algorithm(self): with pytest.raises(ValueError): LSQConfig(scale_algorithm={"method": "smoothquant"}) + def test_local_hessian_activation_error_coupling(self): + cfg = LSQConfig( + scale_algorithm={"method": "local_hessian", "activation_error_coupling": True} + ) + assert cfg.model_dump()["scale_algorithm"]["activation_error_coupling"] is True + def test_scale_algorithm_preserves_sparse_dict(self, monkeypatch): cfg = LSQConfig(scale_algorithm={"method": "mse", "fp8_scale_sweep": True}) assert cfg.model_dump()["scale_algorithm"] == {