diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index ac20729da19..e8888b1789f 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -400,8 +400,8 @@ def max_calibrate( skip_forward_without_activation_calib: If True, skip the (potentially expensive) ``forward_loop`` when no enabled quantizer needs data-driven activation statistics (see :func:`_needs_activation_forward_for_max_calib`). Weight calibration still runs. - Opted into by the top-level ``max`` path and by GPTQ and imatrix, which run their own - forward; algorithms that always need activations (MSE, local Hessian, SmoothQuant, + Opted into by the top-level ``max`` path and by GPTQ, imatrix and local Hessian, which + run their own forward; algorithms that always need activations (MSE, SmoothQuant, SVDQuant) call this with the default False. See :class:`MaxCalibConfig ` for @@ -1214,9 +1214,16 @@ def local_hessian_calibrate( warnings.warn("forward_loop must be provided for local_hessian; skipping local_hessian") return - # Phase 1: max-calibrate (also bootstraps dead experts + promotes/syncs NVFP4 static). + # Phase 1: max-calibrate (also bootstraps dead experts + promotes/syncs NVFP4 static). The + # Hessian pass below is local Hessian's own forward; max only runs one if activations need amax. print_rank_0("local_hessian: Running max calibration for all quantizers...") - max_calibrate(model, forward_loop, distributed_sync, shared_states=shared_states) + max_calibrate( + model, + forward_loop, + distributed_sync, + shared_states=shared_states, + skip_forward_without_activation_calib=True, + ) names = module_name_maps(model) diff --git a/tests/unit/torch/quantization/test_local_hessian.py b/tests/unit/torch/quantization/test_local_hessian.py index d8b4c4eecec..52f53664a9a 100644 --- a/tests/unit/torch/quantization/test_local_hessian.py +++ b/tests/unit/torch/quantization/test_local_hessian.py @@ -132,6 +132,29 @@ 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 + @pytest.mark.parametrize( + ("input_cfg", "forwards"), + [(None, 1), ({"num_bits": 8, "axis": None}, 2)], + ids=["weight_only", "static_input_amax"], + ) + def test_max_calibration_runs_a_forward_only_for_activation_amax(self, input_cfg, forwards): + """Local Hessian's own forward collects its Hessians, so max calibration adds one only + when an activation quantizer needs a data-driven amax.""" + quant_cfg = {**INT8_WEIGHT_CFG, "quant_cfg": list(INT8_WEIGHT_CFG["quant_cfg"])} + if input_cfg is not None: + quant_cfg["quant_cfg"].append({"quantizer_name": "*input_quantizer", "cfg": input_cfg}) + torch.manual_seed(0) + model = SimpleLinear() + mtq.quantize(model, quant_cfg, forward_loop=_make_forward_loop()) + calls = [] + + def forward_loop(model): + calls.append(None) + _make_forward_loop()(model) + + local_hessian_calibrate(model, forward_loop, fp8_scale_sweep=False) + assert len(calls) == forwards + def test_warns_with_module_name_when_cin_not_divisible(self): class _OddModel(nn.Module): def __init__(self):