diff --git a/modelopt/recipe/presets.py b/modelopt/recipe/presets.py index bb4183c790a..5a0d51f2a38 100644 --- a/modelopt/recipe/presets.py +++ b/modelopt/recipe/presets.py @@ -75,6 +75,7 @@ "nvfp4_awq": "nvfp4_awq_lite", "nvfp4_mse": "nvfp4_w4a4_weight_mse_fp8_sweep", "nvfp4_local_hessian": "nvfp4_w4a4_weight_local_hessian", + "nvfp4_local_hessian_act_aware": "nvfp4_w4a4_weight_local_hessian_act_aware", "fp8_pb_wo": "fp8_2d_blockwise_weight_only", "fp8_pc_pt": "fp8_per_channel_per_token", } diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py index a56588b21d8..b6e6c0b1d43 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 + cross_ptr, # [N_CIN_BLOCKS * BLOCK_SIZE * BLOCK_SIZE] fp32 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_CROSS: tl.constexpr, ): pid = tl.program_id(axis=0) cin_block = pid % N_CIN_BLOCKS @@ -205,6 +207,14 @@ def _fp8_scale_sweep_hessian_kernel( + idx[:, None] * BLOCK_SIZE + idx[None, :] ).to(tl.float32) # [BS, BS] + if HAS_CROSS: + cross_t = tl.load( + cross_ptr + + cin_block * (BLOCK_SIZE * BLOCK_SIZE) + + idx[None, :] * BLOCK_SIZE + + idx[:, None] + ).to(tl.float32) # [BS, BS], transposed + pw = tl.dot(w, cross_t, allow_tf32=False) # [ROWS, BS] best_loss = tl.full([ROWS_PER_PROGRAM], float("inf"), dtype=tl.float32) best_idx = tl.zeros([ROWS_PER_PROGRAM], dtype=tl.int32) @@ -218,6 +228,8 @@ def _fp8_scale_sweep_hessian_kernel( # 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_CROSS: + loss -= 2.0 * tl.sum(dw * pw, 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) @@ -230,12 +242,14 @@ def nvfp4_fp8_scale_sweep_hessian( x: torch.Tensor, global_amax: torch.Tensor, hessian: torch.Tensor, + cross: torch.Tensor | None = None, block_size: int = 16, ) -> 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 ``dwᵀ H dw - 2 dwᵀ P w`` (``dw = w - quant(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. @@ -246,6 +260,7 @@ def nvfp4_fp8_scale_sweep_hessian( 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). + cross: Optional activation-aware cross term with the same shape as ``hessian``. block_size: NVFP4 block size (typically 16). Returns: @@ -258,6 +273,11 @@ def nvfp4_fp8_scale_sweep_hessian( f"got {tuple(hessian.shape)}." ) n_cin_blocks = hessian.shape[0] + if cross is not None and cross.shape != hessian.shape: + raise ValueError( + f"cross must have the same shape as hessian {tuple(hessian.shape)}, " + f"got {tuple(cross.shape)}." + ) if n_blocks % n_cin_blocks != 0: raise ValueError( f"n_blocks ({n_blocks}) is not divisible by n_cin_blocks ({n_cin_blocks})." @@ -274,9 +294,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) + cross_flat = ( + cross.contiguous().to(device=x.device, dtype=torch.float32).view(-1) + if cross is not None + else hessian_flat + ) _fp8_scale_sweep_hessian_kernel[grid]( x_flat, hessian_flat, + cross_flat, candidate_scales, candidate_amaxes, best_amax, @@ -285,6 +311,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_CROSS=cross 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..1b8b0ca546b 100644 --- a/modelopt/torch/quantization/calib/mse.py +++ b/modelopt/torch/quantization/calib/mse.py @@ -196,16 +196,18 @@ def __init__( quant_func: Callable | None = None, error_func: Callable | None = None, hessian: torch.Tensor | None = None, + cross: torch.Tensor | None = None, ): """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]``) and its + optional activation-aware ``cross`` term enable the local-Hessian Triton fast path; + ``error_func`` carries the same metric for the reference fallback. """ 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._cross = cross # Set by collect() after either sweep path; consumed by compute_amax. self._best_amax: torch.Tensor | None = None @@ -265,7 +267,11 @@ def collect(self, x: torch.Tensor): from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep_hessian best_flat = nvfp4_fp8_scale_sweep_hessian( - x.detach(), self._global_amax, self._hessian, block_size=x.shape[-1] + x.detach(), + self._global_amax, + self._hessian, + cross=self._cross, + block_size=x.shape[-1], ) 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 a8ed323a7b2..1acd2d15a2a 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -995,7 +995,8 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): 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. + with the local Hessian matrix computed from input activations. For best calibration + fidelity and bounded memory, use ``batch_size=1`` and ``layerwise=True``. The local Hessian loss for each block is: ``(dw @ H @ dw.T)`` where: - ``dw = weight - quantized_weight`` (weight reconstruction error per block) @@ -1044,6 +1045,13 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): "Default is 16 for NVFP4.", ) + act_quant_aware: bool = ModeloptField( + default=False, + title="Include activation quantization in the scale-search objective.", + description="If True, minimize the exact expanded output error for quantized weights " + "and activations using H = XqᵀXq and P = Xqᵀ(Xq - X).", + ) + distributed_sync: bool | None = ModeloptField( default=True, title="Whether to sync the amax across the distributed processes.", @@ -1652,6 +1660,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_AWARE_CFG: dict[str, Any] = _load_quantize_config_dict( + "configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_aware" +) MAMBA_MOE_NVFP4_AGGRESSIVE_CFG: dict[str, Any] = _load_quantize_config_dict( "configs/ptq/presets/model/mamba_moe_nvfp4_aggressive" ) @@ -1743,6 +1754,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_AWARE_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 e03dff2e98b..02b69b30ae9 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -465,11 +465,12 @@ def _make_weight_mse_calibrator( fp8_scale_sweep: bool, error_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, hessian: torch.Tensor | None = None, + cross: torch.Tensor | 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 + ``hessian`` and optional activation-aware ``cross`` tensors enable NVFP4's Hessian-weighted Triton fast path; ``error_func`` then serves only as the reference fallback. """ if ( @@ -507,6 +508,7 @@ def _make_weight_mse_calibrator( quant_func=quant_func, error_func=error_func, hessian=hessian, + cross=cross, ) # fp8_scale_sweep applies only to registered backends and static NVFP4; skip others. return None @@ -580,13 +582,15 @@ def _mse_calibrate_weights( fp8_scale_sweep: bool, error_func_for: Callable[[TensorQuantizer], Callable | None] | None = None, hessian_for: Callable[[TensorQuantizer], torch.Tensor | None] | None = None, + cross_for: Callable[[TensorQuantizer], torch.Tensor | 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. + enabling the Hessian-weighted Triton fast path. ``cross_for`` supplies its optional + activation-quantization-aware cross term. """ seen_modules: set[int] = set() pbar = tqdm(desc="MSE weight calibration") @@ -598,6 +602,7 @@ def _mse_calibrate_weights( 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 + cross = cross_for(weight_quantizer) if cross_for else None cal = _make_weight_mse_calibrator( weight_quantizer, step_size, @@ -606,6 +611,7 @@ def _mse_calibrate_weights( fp8_scale_sweep, error_func=error_func, hessian=hessian, + cross=cross, ) if cal is None: continue @@ -621,7 +627,7 @@ def _mse_calibrate_weights( class _LocalHessianAccumulator: - """Per-block local Hessian ``H = ΣXᵀX`` for one weight quantizer. + """Per-block local Hessian and activation-aware cross term 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. @@ -635,22 +641,37 @@ 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.cross_per_block: torch.Tensor | None = None self._normalized_hessian: torch.Tensor | None = None + self._normalized_cross: 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, input_tensor: torch.Tensor, quantized_input: torch.Tensor | None = None + ) -> None: + """Accumulate per-block Hessian and optional activation-aware cross term.""" 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 = x.reshape(self.num_blocks_per_cin, self.block_size, -1) - hessian_batch = x @ x.transpose(-1, -2) + if quantized_input is None: + xq = x + else: + xq = quantized_input.reshape(-1, self.cin).to(torch.float32).T + xq = xq.reshape(self.num_blocks_per_cin, self.block_size, -1) + hessian_batch = xq @ xq.transpose(-1, -2) if self.hessian_per_block is None: self.hessian_per_block = hessian_batch else: self.hessian_per_block += hessian_batch + if quantized_input is not None: + cross_batch = xq @ (xq - x).transpose(-1, -2) + if self.cross_per_block is None: + self.cross_per_block = cross_batch + else: + self.cross_per_block += cross_batch self.num_samples += input_tensor.numel() // self.cin def normalized_hessian(self) -> torch.Tensor | None: @@ -667,26 +688,37 @@ def normalized_hessian(self) -> torch.Tensor | None: self._normalized_hessian = self.hessian_per_block / self.num_samples return self._normalized_hessian + def normalized_cross(self) -> torch.Tensor | None: + """Per-cin-block cross term ``Xqᵀ(Xq - X) / num_samples`` if captured.""" + if self._normalized_cross is None and self.cross_per_block is not None and self.num_samples: + self._normalized_cross = self.cross_per_block / self.num_samples + return self._normalized_cross + def build_error_func( self, keep_buffer: bool = False ) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None: """Hessian-weighted error function (``None`` if no samples). - Frees the raw Hessian buffer unless ``keep_buffer`` (kept for debug inspection). + Frees the raw Hessian and cross buffers unless ``keep_buffer`` (for debug inspection). """ hessian = self.normalized_hessian() + cross = self.normalized_cross() if hessian is None: return None cout = self.cout bs = self.block_size if not keep_buffer: self.hessian_per_block = None + self.cross_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) + block_loss = torch.einsum("cnb,nbd,cnd->cn", dw, hessian, dw) + if cross is not None: + block_loss -= 2 * torch.einsum("cnb,nbd,cnd->cn", dw, cross, x.view(cout, -1, bs)) + block_loss = block_loss.reshape(-1) return block_loss.unsqueeze(-1).expand(-1, bs).reshape(original_shape) return local_hessian_error @@ -727,7 +759,9 @@ def _is_quant_fused_experts(module: nn.Module) -> bool: ) -def _register_local_hessian_input_hooks(model, name_to_module, capture, block_size, warned): +def _register_local_hessian_input_hooks( + model, name_to_module, capture, block_size, warned, act_quant_aware=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 @@ -737,14 +771,32 @@ 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 _can_capture_quantized_input(input_quantizer): + return ( + isinstance(input_quantizer, TensorQuantizer) + and input_quantizer.pre_quant_scale is None + and not input_quantizer.rotate_is_enabled + ) + + def _warn_act_quant_fallback(name): + warn_rank_0( + f"local_hessian: {name} input quantizer has pre_quant_scale or rotation, or is not " + "a TensorQuantizer; falling back to non-activation-aware capture for this layer." + ) + + def _make_expert_hook(expert_module, weight_name, quantizers, enabled, capture_quantized): + def _expert_hook(_input_quantizer, args, output=None): 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], + args[0], + output if capture_quantized else None, + ) return _expert_hook @@ -760,11 +812,25 @@ def _expert_hook(_input_quantizer, args): name, weight, module.weight_quantizer, block_size, warned ) + input_quantizer = getattr(module, "input_quantizer", None) + capture_quantized = act_quant_aware and _can_capture_quantized_input(input_quantizer) + if act_quant_aware and not capture_quantized: + _warn_act_quant_fallback(name) + def _dense_hook(linear, args): if args: capture(linear.weight_quantizer, linear.weight, args[0]) - handles.append(module.register_forward_pre_hook(_dense_hook)) + if capture_quantized: + assert isinstance(input_quantizer, TensorQuantizer) + + def _dense_quantizer_hook(_input_quantizer, args, output, linear=module): + if args: + capture(linear.weight_quantizer, linear.weight, args[0], output) + + handles.append(input_quantizer.register_forward_hook(_dense_quantizer_hook)) + else: + handles.append(module.register_forward_pre_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") @@ -787,10 +853,18 @@ 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} + capture_quantized = act_quant_aware and _can_capture_quantized_input( + input_quantizer + ) + if act_quant_aware and not capture_quantized: + _warn_act_quant_fallback(f"{name}.{weight_name}") + hook = _make_expert_hook( + module, weight_name, quantizers, enabled, capture_quantized + ) handles.append( - input_quantizer.register_forward_pre_hook( - _make_expert_hook(module, weight_name, quantizers, enabled) - ) + input_quantizer.register_forward_hook(hook) + if capture_quantized + else input_quantizer.register_forward_pre_hook(hook) ) return handles @@ -805,6 +879,7 @@ def local_hessian_calibrate( stop_multiplier: float = 4.0, fp8_scale_sweep: bool = True, block_size: int = 16, + act_quant_aware: bool = False, debug: bool = False, shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, ): @@ -812,7 +887,9 @@ def local_hessian_calibrate( 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``. + (input quantizers untouched) and fed to :func:`mse_calibrate`'s weight search via + ``error_func``. With ``act_quant_aware=True``, it instead minimizes the exact second-order + expansion of ``||Xq Wq - X W||²`` using Hessian ``XqᵀXq`` and cross term ``Xqᵀ(Xq - X)``. 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 @@ -830,6 +907,7 @@ 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). + act_quant_aware: Include activation quantization in the local objective (default: False). debug: If True, retain the per-quantizer Hessian accumulators on the model (``model._local_hessian_accumulators``) for inspection. @@ -849,19 +927,23 @@ 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): + def capture(weight_quantizer, weight, input_tensor, quantized_input=None): input_local = input_tensor.to_local() if hasattr(input_tensor, "to_local") else input_tensor + quantized_local = ( + quantized_input.to_local() if hasattr(quantized_input, "to_local") else quantized_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(input_local, quantized_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. + # disabled. Input quantizers are left as-is, yielding H = ΣXᵀX normally or H = ΣXqᵀXq + # plus the activation-aware cross term when requested. 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, act_quant_aware ) print_rank_0("local_hessian: Caching activations and computing local Hessian...") try: @@ -888,6 +970,7 @@ def capture(weight_quantizer, weight, input_tensor): 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()} + crosses = {qid: acc.normalized_cross() for qid, acc in accumulators.items()} print_rank_0("local_hessian: Running MSE calibration with local Hessian loss...") _mse_calibrate_weights( model, @@ -898,17 +981,20 @@ def capture(weight_quantizer, weight, input_tensor): fp8_scale_sweep=fp8_scale_sweep, error_func_for=lambda q: error_funcs.get(id(q)), hessian_for=lambda q: hessians.get(id(q)), + cross_for=lambda q: crosses.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. + # Release the per-block metrics (held by the error_func closures, calibrators, and the + # accumulators' caches) before empty_cache so export starts defragmented; keep only for debug. error_funcs.clear() hessians.clear() + crosses.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._cross = None if debug: model._local_hessian_accumulators = accumulators else: diff --git a/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_aware.yaml b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_aware.yaml new file mode 100644 index 00000000000..610c015e80b --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/nvfp4_w4a4_weight_local_hessian_act_aware.yaml @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 activation-aware local-Hessian weight scales. + +# 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 + act_quant_aware: 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..275b82f7251 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, hessian_input, quantized_input=None, block_size=BLOCK_SIZE +): """Real ``_LocalHessianAccumulator`` so the test exercises the production metric.""" acc = _LocalHessianAccumulator(cout, cin, block_size) - acc.accumulate(hessian_input) + acc.accumulate(hessian_input, quantized_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(), + cross=acc.normalized_cross(), ) cal.collect(x_blocks) return cal.compute_amax() @@ -428,6 +431,18 @@ def _total_hessian_loss(x_blocks, per_block_amax, global_amax, hessian): return (torch.einsum("nij,nj->ni", h_per_block, dw) * dw).sum() +def _activation_aware_loss(x_blocks, per_block_amax, global_amax, hessian, cross): + """Per-block activation-aware objective at the selected weight scales.""" + n_blocks = x_blocks.shape[0] + n_cin = hessian.shape[0] + cin_idx = 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 + quad = (torch.einsum("nij,nj->ni", hessian[cin_idx], dw) * dw).sum(dim=1) + pw = torch.einsum("nij,nj->ni", cross[cin_idx], x_blocks.float()) + return quad - 2 * (dw * pw).sum(dim=1) + + @requires_triton @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize( @@ -485,6 +500,35 @@ def test_hessian_parity_random_weights(seed, cout, cin, dtype): ) +@requires_triton +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_activation_aware_hessian_parity(dtype): + """Activation-aware Hessian+cross Triton sweep matches the reference objective.""" + torch.manual_seed(17) + cout, cin = 32, 64 + weight = torch.randn(cout, cin, device="cuda", dtype=dtype) + hessian_input = torch.randn(512, cin, device="cuda") + # A systematic scale bias makes weight compensation reduce the activation-only error, + # deterministically exercising negative values of the constant-dropped objective. + quantized_input = (torch.round(hessian_input * 2) / 2) * 1.25 + acc = _build_hessian_accumulator(cout, cin, hessian_input, quantized_input) + 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) + n_diff = int((ref != tri).sum()) + if dtype != torch.bfloat16: + assert torch.equal(ref, tri), f"{n_diff}/{ref.numel()} blocks differ" + else: + assert n_diff / ref.numel() < 1e-3 + winning_loss = _activation_aware_loss( + x_blocks, tri, global_amax, acc.normalized_hessian(), acc.normalized_cross() + ) + assert (winning_loss < 0).any() + + @requires_triton def test_hessian_sweep_input_validation(): """``nvfp4_fp8_scale_sweep_hessian`` should reject malformed inputs cleanly.""" @@ -501,6 +545,58 @@ 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="cross must have the same shape"): + nvfp4_fp8_scale_sweep_hessian(x, g, h, cross=torch.randn(4, 8, 8, device=device)) + + +@requires_triton +def test_activation_aware_local_hessian_end_to_end(monkeypatch): + """Activation-aware local-Hessian fast and reference paths select identical amaxes.""" + if get_cuda_ext_mx() is None: + pytest.skip("cuda_ext_mx is not available") + cfg = { + "quant_cfg": [ + { + "quantizer_name": "*weight_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)}, + "axis": None, + }, + "enable": True, + }, + { + "quantizer_name": "*input_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + }, + "enable": True, + }, + ], + "algorithm": { + "method": "local_hessian", + "fp8_scale_sweep": True, + "act_quant_aware": True, + }, + } + + def _run(env_value): + torch.manual_seed(19) + model = SimpleLinear().cuda() + monkeypatch.setenv("MODELOPT_NVFP4_TRITON_SWEEP", env_value) + data = [model.get_input().cuda() for _ in range(2)] + mtq.quantize(model, cfg, forward_loop=lambda m: [m(x) for x in data]) + return { + name: module.amax.clone() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.is_nvfp4_static + } + + fast = _run("1") + reference = _run("0") + assert fast.keys() == reference.keys() + assert all(torch.equal(fast[name], reference[name]) for name in fast) @requires_triton diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 9fa836bb620..6e454cfa04e 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("act_quant_aware", [False, True]) + def test_local_hessian_refines_per_expert_weights(self, act_quant_aware): """local_hessian captures each expert's routed activations and refines its weight amax.""" model = _TinyMoEModel() expert_type = type(model.moe.experts) @@ -936,7 +937,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, + act_quant_aware=act_quant_aware, + 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 +951,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.normalized_cross() is not None) == act_quant_aware 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..3c2027abd17 100644 --- a/tests/unit/torch/quantization/test_local_hessian.py +++ b/tests/unit/torch/quantization/test_local_hessian.py @@ -49,6 +49,15 @@ "algorithm": "max", } +INT8_WEIGHT_ACT_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 { @@ -81,6 +90,56 @@ def test_accumulate_shape_samples_fp32_buffer(self): assert acc.num_samples == 15 assert acc.build_error_func() is not None assert acc.hessian_per_block is None # raw buffer freed + assert acc.cross_per_block is None # default path never allocates the cross term + + def test_activation_aware_accumulator_matches_explicit_einsums(self): + torch.manual_seed(2) + cin, bs = 32, 16 + acc = _LocalHessianAccumulator(4, cin, bs) + x = torch.randn(7, cin) + xq = x + 0.1 * torch.randn_like(x) + acc.accumulate(x, xq) + + xb = x.T.reshape(cin // bs, bs, -1) + xqb = xq.T.reshape(cin // bs, bs, -1) + expected_hessian = torch.einsum("nbt,ndt->nbd", xqb, xqb) / x.shape[0] + expected_cross = torch.einsum("nbt,ndt->nbd", xqb, xqb - xb) / x.shape[0] + assert torch.allclose(acc.normalized_hessian(), expected_hessian) + assert torch.allclose(acc.normalized_cross(), expected_cross) + + assert acc.build_error_func() is not None + assert acc.hessian_per_block is None + assert acc.cross_per_block is None + + def test_activation_aware_error_matches_exact_objective_and_pins_sign(self): + torch.manual_seed(3) + tokens, cout, cin = 11, 3, 16 + x = torch.randn(tokens, cin) + xq = x + 0.2 * torch.randn_like(x) + w0 = torch.randn(cout, cin) + wq = w0 + 0.1 * torch.randn_like(w0) + acc = _LocalHessianAccumulator(cout, cin, cin) + acc.accumulate(x, xq) + + actual = acc.build_error_func()(w0, wq)[:, 0] + full_error = (xq @ wq.T - x @ w0.T).square().sum(dim=0) + activation_constant = (xq @ w0.T - x @ w0.T).square().sum(dim=0) + expected = (full_error - activation_constant) / tokens + assert torch.allclose(actual, expected, atol=1e-5) + + def test_zero_cross_reduces_to_existing_objective(self): + torch.manual_seed(4) + x = torch.randn(9, 16) + w = torch.randn(2, 16) + wq = w + 0.1 * torch.randn_like(w) + plain = _LocalHessianAccumulator(2, 16, 16) + aware = _LocalHessianAccumulator(2, 16, 16) + plain.accumulate(x) + aware.accumulate(x, x) + + assert torch.count_nonzero(aware.normalized_cross()) == 0 + assert torch.equal(plain.normalized_hessian(), aware.normalized_hessian()) + assert torch.equal(plain.build_error_func()(w, wq), aware.build_error_func()(w, wq)) def test_error_func_matches_explicit_hessian_weighted_loss(self): torch.manual_seed(1) @@ -132,6 +191,67 @@ def test_refines_amax_beyond_max_and_plain_mse(self): assert any(not torch.allclose(lh[n], max_amax[n]) for n in lh) # refined past max-cal assert any(not torch.allclose(lh[n], mse[n]) for n in lh) # Hessian changed the choice + def test_activation_aware_capture_and_disabled_quantizer_degeneracy(self): + forward_loop = _make_forward_loop() + torch.manual_seed(5) + aware_model = SimpleLinear() + mtq.quantize(aware_model, INT8_WEIGHT_ACT_CFG, forward_loop=forward_loop) + local_hessian_calibrate( + aware_model, + forward_loop, + fp8_scale_sweep=False, + act_quant_aware=True, + debug=True, + ) + assert any( + acc.normalized_cross() is not None + for acc in aware_model._local_hessian_accumulators.values() + ) + torch.manual_seed(5) + plain_model = SimpleLinear() + mtq.quantize(plain_model, INT8_WEIGHT_ACT_CFG, forward_loop=forward_loop) + local_hessian_calibrate(plain_model, forward_loop, fp8_scale_sweep=False) + aware_amaxes = _weight_amaxes(aware_model) + plain_amaxes = _weight_amaxes(plain_model) + assert any( + not torch.allclose(aware_amaxes[name], plain_amaxes[name]) for name in aware_amaxes + ) + + def _run_disabled(act_quant_aware): + torch.manual_seed(6) + model = SimpleLinear() + mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=forward_loop) + local_hessian_calibrate( + model, + forward_loop, + fp8_scale_sweep=False, + act_quant_aware=act_quant_aware, + ) + return _weight_amaxes(model) + + plain = _run_disabled(False) + aware_disabled = _run_disabled(True) + assert plain.keys() == aware_disabled.keys() + assert all(torch.equal(plain[name], aware_disabled[name]) for name in plain) + + def test_activation_aware_falls_back_for_pre_quant_scale(self): + forward_loop = _make_forward_loop() + torch.manual_seed(7) + model = SimpleLinear() + mtq.quantize(model, INT8_WEIGHT_ACT_CFG, forward_loop=forward_loop) + linear = model.net[0] + linear.input_quantizer.pre_quant_scale = torch.ones(16) + with pytest.warns(UserWarning, match="falling back to non-activation-aware capture"): + local_hessian_calibrate( + model, + forward_loop, + fp8_scale_sweep=False, + act_quant_aware=True, + debug=True, + ) + acc = model._local_hessian_accumulators[id(linear.weight_quantizer)] + assert acc.normalized_cross() is None + def test_warns_with_module_name_when_cin_not_divisible(self): class _OddModel(nn.Module): def __init__(self):