model_calib#

Calibration utilities.

Functions

`awq`	Apply AWQ to the model.
`layerwise_calibrate`	Layerwise calibration - a layer-by-layer calibration algorithm.
`local_hessian_calibrate`	Calibrate weight quantizers by minimizing the Hessian-weighted error.
`lsq`	Run scale calibration then convert to LSQ mode.
`max_calibrate`	Calibrate the model using max.
`mse_calibrate`	Calibrate weight quantizers using MSE-based amax search.
`nvfp4_act_headroom_calibrate`	Calibrate NVFP4 activation global scales with headroom.
`smoothquant`	Smooth-Quant variant with per-channel weight scaling.
`svdquant`	Lite version of SVDQuant.

awq(model, forward_loop=None, algorithm='awq_lite', **kwargs)#

Apply AWQ to the model.

Parameters:

model (Module) – Model to be calibrated.
forward_loop (Callable[[Module], None] | None) – A callable which takes the model as argument and forwards calibration data through the model.
algorithm (str)

See AWQFullCalibConfig for details on the remaining arguments.

layerwise_calibrate(model, forward_loop, calib_func, **calib_kwargs)#

Layerwise calibration - a layer-by-layer calibration algorithm.

Runs the full model forward per layer but patches decoder layers with a skip / run / capture strategy so that inter-layer logic in parent modules (e.g. mask construction) executes naturally without model-specific hooks.

Every knob arrives through calib_kwargs from LayerwiseConfig, which documents them; export_dir additionally leaves the model in export form, so it must not be used for inference afterwards.

Parameters:

model (Module)
forward_loop (Callable[[Module], None])
calib_func (Callable)

local_hessian_calibrate(model, forward_loop=None, distributed_sync=True, step_size=0.1, start_multiplier=0.25, stop_multiplier=4.0, fp8_scale_sweep=True, block_size=16, debug=False, shared_states=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 mse_calibrate()’s weight search via error_func.

Like 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.

Parameters:

model (Module) – Model to be calibrated.
forward_loop (Callable[[Module], None] | None) – A callable which takes the model as argument and forwards calibration data through the model. Required for this algorithm.
distributed_sync (bool) – Whether to sync amax across distributed processes.
step_size (float) – Step size for amax search (default: 0.1).
start_multiplier (float) – Starting multiplier for amax search (default: 0.25).
stop_multiplier (float) – Ending multiplier for amax search (default: 4.0).
fp8_scale_sweep (bool) – If True, sweep over all 128 possible FP8 E4M3 scale values for NVFP4 per-block quantization (default: True).
block_size (int) – Block size for local Hessian computation (default: 16).
debug (bool) – If True, retain the per-quantizer Hessian accumulators on the model (model._local_hessian_accumulators) for inspection.
shared_states (Mapping[str, Mapping[str, Sequence[str]]] | None)

See LocalHessianCalibConfig for details on the configuration options.

lsq(model, forward_loop=None, scale_algorithm=None, learnable_amax=('post',), tied_amax=False, quantize_pre_scale=True, **kwargs)#

Run scale calibration then convert to LSQ mode.

Uses separate pre (quant) and post (dequant) amax values. Forward: w_q = Q_STE(w / s_pre) * s_post where s = amax / Q_max.

Parameters:

model (Module) – Quantized model.
forward_loop (Callable[[Module], None] | None) – Calibration data forward loop.
scale_algorithm (dict | None) – Calibration algorithm config to run first. Dict with ‘method’ key: ‘mse’, ‘local_hessian’, or ‘max’. Defaults to {‘method’: ‘mse’} if None.
learnable_amax (list | str) – Which amax params are learnable: ‘pre’, ‘post’, [‘pre’, ‘post’], or [].
tied_amax (bool) – If True, pre and post share a single tensor.
quantize_pre_scale (bool) – If False, skip FP8 quantization for the LSQ pre scale.

max_calibrate(model, forward_loop=None, distributed_sync=True, sync_expert_weight_amax=False, shared_states=None, skip_forward_without_activation_calib=False)#

Calibrate the model using max.

Parameters:

model (Module) – Model to be calibrated.
forward_loop (Callable[[Module], None] | None) – A callable which takes the model as argument and forwards calibration data through the model.
distributed_sync – Whether to sync input_quantizer amax across distributed processes.
sync_expert_weight_amax – SequentialMLP only — share one weight amax across all experts in a MoE layer (within-rank sync + EP all-reduce when EP>1).
shared_states (Mapping[str, Mapping[str, Sequence[str]]] | None) – Optional dict keyed by shared-state name. "weight_global_amax" is implemented today and accepts {"patterns": [...]}; omitted patterns use SHARED_PATTERNS, while an empty list disables the state.
skip_forward_without_activation_calib (bool) – If True, skip the (potentially expensive) forward_loop when no enabled quantizer needs data-driven activation statistics (see _needs_activation_forward_for_max_calib()). Weight calibration still runs. Only opt-in for the top-level max path; algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call this with the default False.

See MaxCalibConfig for details on the remaining arguments.

mse_calibrate(model, forward_loop=None, distributed_sync=True, step_size=0.1, start_multiplier=0.25, stop_multiplier=4.0, fp8_scale_sweep=False, shared_states=None)#

Calibrate weight quantizers using MSE-based amax search.

This calibration method first uses max calibration to initialize amax values for all quantizers, then searches for better weight amax values by minimizing the MSE between original and quantized weights.

Parameters:

model (Module) – Model to be calibrated.
forward_loop (Callable[[Module], None] | None) – A callable which takes the model as argument and forwards calibration data through the model.
distributed_sync – Whether to sync amax across distributed processes.
step_size (float) – Step size for amax search (default: 0.1).
start_multiplier (float) – Starting multiplier for amax search (default: 0.25).
stop_multiplier (float) – Ending multiplier for amax search (default: 4.0).
fp8_scale_sweep (bool) – If True, only ModelOpt static NVFP4 weights and registered custom backends are MSE-calibrated (via FP8 E4M3 scale-value sweep); all other weight quantizers (INT8, plain FP8, unregistered backends, etc.) are skipped and left at their max-calibrated amax. If False, all weight quantizers use the multiplier search.
shared_states (Mapping[str, Mapping[str, Sequence[str]]] | None)

See MseCalibConfig for details on the remaining arguments.

nvfp4_act_headroom_calibrate(model, forward_loop=None, *, anchor_percentile=1.0, upper_percentile=99.99, rho=16384.0, weight_scale_algorithm=None)#

Calibrate NVFP4 activation global scales with headroom.

For NVFP4 dynamic-block input quantizers, the per-tensor global scale is derived from the distribution of per-block activation amaxes so that headroom is left above the calibrated range (see NVFP4ActHeadroomCalibrator).

Weight scales are an orthogonal concern and are delegated to weight_scale_algorithm, so a recipe can combine this activation policy with max, mse or local_hessian weights. Each of those runs max calibration first, so the activation collectors installed here are populated in that pass and any later refinement touches only weights.

Parameters:

model (Module) – model to be calibrated.
forward_loop (Callable[[Module], None] | None) – callable that runs calibration data through the model.
anchor_percentile (float) – percentile of the per-block amaxes used as the anchor.
upper_percentile (float) – percentile of the per-block amaxes used as the top of the calibrated range; 100 uses the literal observed max.
rho (float) – headroom factor; amax = rho * anchor. Must be in (0, 28672).
weight_scale_algorithm (Any) – config for the algorithm that calibrates the weight scales – {"method": "max"} (the default), "mse" or "local_hessian", plus that algorithm’s own options such as distributed_sync and shared_states. None is treated as {"method": "max"}.

Raises:

NotImplementedError – if an NVFP4 activation quantizer is a SequentialQuantizer, which this algorithm does not support.

Note

Under data parallelism the per-rank scales are combined with a MAX all-reduce, so the result is the largest per-rank headroom scale rather than the scale implied by pooling every rank’s per-block distribution.

smoothquant(model, forward_loop=None, alpha=1.0)#

Smooth-Quant variant with per-channel weight scaling.

Parameters:

model (Module) – Model to be calibrated.
forward_loop (Callable[[Module], None] | None) – A callable which takes the model as argument and forwards calibration data through the model.

See SmoothQuantCalibConfig for details on the remaining arguments.

svdquant(model, forward_loop=None, lowrank=32, skip_layers=None, **kwargs)#

Lite version of SVDQuant.

Parameters:

model (Module) – Model to be calibrated.
forward_loop (Callable[[Module], None] | None) – A callable which takes the model as argument and forwards calibration data through the model.
lowrank (int)
skip_layers (list[str] | None)

See SVDQuantConfig for details on the remaining arguments.