Skip to content

ParoQuant

Jump to bottom
Qubitium edited this page Nov 2, 2026 · 5 revisions

ParoQuant

Documentation home · Quickstart · Lifecycle

On this page

Overview

ParoQuant quantizes linear weights in a learned, pairwise-rotated coordinate system. Calibration optimizes Givens-rotation angles and channel scales, then fine-tunes weights and quantizer parameters against floating-point reconstruction targets. Inference rotates activations and multiplies them by packed, transformed-domain weights.

The method is METHOD.PARO ("paroquant"); its checkpoint format is FORMAT.PAROQUANT ("paroquant"). It is not an ordinary AWQ checkpoint, even though its integer weights use AWQ packing: the rotations and channel scales are part of the model and must be loaded with the weights.

This reference describes the implementation at source revision eb90dbdd5416c0c8e50d1b64f1ac562012afe4aa. Backend support and optimization controls below refer to that revision.

Workflow at a glance

Learned coordinates and packed inference. ParoQuant stores transformed weights together with required rotation metadata.

Reading the diagram: Independent optimization and execution concerns. Grouped scopes use layer replay objectives; zero epochs skip updates, not export. Packing is not a resume checkpoint.

Text version of the workflow
  • Module-scope optimization: Capture activation rows → Learn angles and scales → Fine-tune quantization → Pack runtime state.
  • Packed inference: Scale input channels → Apply learned rotations → Evaluate packed weights → Restore output shape.

Open full-size SVG

Concepts

Pairwise Givens rotations

For every group of input channels, a rotation stage divides channels into disjoint pairs. For one pair with values $a,b$ and angle $\theta$, the reference implementation computes:

$$ a'=a\cos\theta+b\sin\theta,\qquad b'=-a\sin\theta+b\cos\theta. $$

Stages are applied in order. The inverse reverses stage order and negates the angles. Pair indices are local to each group, not global feature indices. Groups must have even width, and the input width must be divisible by that width. group_size=-1 is normalized to the complete input width.

krot is the number of stages. opt_pair_ratio determines the number of active disjoint pairs per group; unused channels are filled with dummy pairs whose angles stay zero. The fast schedule uses randomized permutations of a round-robin pairing template; the reference schedule uses its separate pair-selection/alignment implementation. These are selectable implementations, not a guarantee of identical random schedules.

Activation and weight transforms

Let $R$ denote the combined orthogonal row-vector rotation and $D$ the diagonal matrix of bounded, learned channel scales. For row-vector activations $X$ and weights $W$ in (out_features, in_features) order, the optimizer forms:

$$ X_r=XD^{-1}R,\qquad W_r=WDR. $$

Without quantization, $X_rW_r^T=XW^T$. With transformed-domain pseudo-quantization $Q$, the replayed linear is:

$$ Y=X_rQ(W_r)^T+b. $$

The equivalent weight in the original input coordinate system is $Q(W_r)R^TD^{-1}$. The processor uses this pseudo-weight during calibration replay, but packs Q(W_r) for runtime. The saved channel_scales are the reciprocal runtime scales $D^{-1}$, not the unconstrained optimization parameter channel_scales_opt.

Straight-through estimation

_RoundSTE rounds in the forward pass and passes its gradient through unchanged. _ClampSTE clamps forward values but also passes gradients through. Consequently, optimization can update weights, angles, channel scales, and learned quantizer parameters despite rounding and clipping.

Two quantizer implementations must be distinguished:

opt_quantizer_impl Internal transformed-weight quantizer
"reference" (configuration default) _quantizer_sym_for_impl returns False: affine quantization with a floating zero-point parameter.
"fast" Preserves sym; with the required sym=True, uses signed symmetric quantization.

The public ParoQuant configuration requires sym=True; setting sym=False is not how to select the affine reference path. Four-bit symmetric quantization uses signed codes from -8 to 7, then exports an unsigned packed representation with zero point 8. Affine quantization uses unsigned codes from 0 to 15. Quantizer scales are bounded to [1e-5, 1e5]; these bounds are separate from the configurable channel-scale bounds.

Usage

Quantize, save, and reload

Use representative calibration text or the supported tokenized calibration inputs described in Quantization lifecycle. This example deliberately supplies the important choices rather than relying on unrelated base-configuration defaults:

from datasets import load_dataset
from gptqmodel import GPTQModel, QuantizeConfig
from gptqmodel.quantization import FORMAT, METHOD

calibration = load_dataset(
    "allenai/c4",
    data_files="en/c4-train.00001-of-01024.json.gz",
    split="train",
).select(range(128))["text"]

qcfg = QuantizeConfig(
    method=METHOD.PARO,
    format=FORMAT.PAROQUANT,
    bits=4,
    group_size=128,
    sym=True,
    krot=8,
    opt_scope="module",
    opt_quantizer_impl="reference",
)
model = GPTQModel.load("meta-llama/Llama-3.2-1B", quantize_config=qcfg)
model.quantize(calibration, batch_size=1)
model.save("llama-paroquant")
from gptqmodel import BACKEND, GPTQModel

model = GPTQModel.load("llama-paroquant", backend=BACKEND.PAROQUANT_CUDA)
tokens = model.generate("Capital of France is", max_new_tokens=32)[0]
print(model.tokenizer.decode(tokens))

For the Triton runtime, request BACKEND.PAROQUANT_TRITON. On supported Apple Silicon systems, use the MLX backend described in MLX on Apple Silicon. Backend selection is independent of opt_stage_impl, opt_pair_impl, and opt_quantizer_impl: those choose calibration implementations, not the inference backend.

Choose optimization scope

opt_scope Objective and orchestration Gradient checkpointing with opt_gradient_checkpointing=None
"module" Each linear independently reconstructs its own floating-point output. Off.
"compute_block" Optimize selected linears together inside a replayed layer against preserved layer outputs. Off.
"layer" Optimize the selected linears of an entire layer together against preserved layer outputs. On.

Compute-block buckets are attn_qkv (q_proj, k_proj, v_proj), attn_o (o_proj), mlp_gate_up (gate_proj, up_proj), and mlp_down (down_proj). Other names receive individual single:<module_name> buckets. These names are implementation buckets, not additional opt_scope values.

opt_gradient_checkpointing trades recomputation for activation memory during optimization. It does not save progress for resuming a run, and neither a best-state snapshot nor model.save() is a quantization continuation checkpoint. See Quantization Lifecycle for the separate continuation API.

qcfg = QuantizeConfig(
    method=METHOD.PARO,
    bits=4,
    group_size=128,
    sym=True,
    opt_scope="layer",
    opt_gradient_checkpointing=True,
    opt_best_state_dtype="bf16",
)

A lower-precision best-state snapshot reduces snapshot storage but rounds the state that will be restored. It does not change the packed bit width. opt_train_samples and opt_validation_samples count flattened activation/token rows, not dataset records. In module scope, opt_batch_size splits these rows into optimization minibatches independently of model.quantize(..., batch_size=...).

Grouped scopes retain complete cached layer-forward batches. Training uses the smallest nonempty prefix meeting the training-row budget; validation uses the smallest nonempty suffix meeting the validation-row budget. These selections can exceed the budgets and can overlap. Grouped runners do not split batches using opt_batch_size: their optimizer steps follow the captured forward batches, so the calibration forward batch size and sequence length directly affect optimization memory.

Configuration reference

ParoConfig is defined in gptqmodel/quantization/config.py and inherits the shared preprocessor/quantization configuration. See QuantizeConfig reference for shared fields such as bits, group_size, dynamic, device selection, and offload settings. The following are the fields declared specifically by ParoConfig.

Method, sampling, and optimization

Field Declared type Default Meaning / validation
method METHOD METHOD.PARO ParoQuant method.
format FORMAT FORMAT.PAROQUANT Only ParoQuant export format; ParoConfig.__post_init__ rewrites another recognized format to this value.
krot int 8 Positive rotation-stage count.
opt_rotation_epochs int 10 Non-negative rotation/channel-scale stage epochs.
opt_finetune_epochs int 10 Non-negative weight/quantizer stage epochs.
opt_train_samples int 2048 Positive training-row budget.
opt_validation_samples int 64 Positive validation-row budget.
opt_batch_size int 64 Positive row minibatch size for module-scope optimization; grouped runners use captured layer batches instead.
opt_rotation_lr float 0.05 Positive angle and channel-scale learning rate.
opt_weight_lr float 1e-5 Positive weight learning rate.
opt_quantizer_lr float 1e-6 Positive learned-quantizer learning rate.
opt_pair_ratio float 0.5 Must satisfy 0 < value <= 0.5.
opt_seed int 0 Base seed for deterministic module-specific schedules.
opt_scope str "module" "module", "compute_block", "layer".
opt_train_on_noisy_inputs bool False Grouped scopes may train on propagated quantized/noisy inputs against clean floating-point targets. Not enabled for module scope.

Optimizer

Field Declared type Default Meaning / validation
opt_optimizer str "adamw" "adamw", "adam", or "sgd".
opt_weight_decay float 0.01 Non-negative.
opt_betas Tuple[float, float] (0.9, 0.95) Two Adam-family coefficients, each in [0, 1).
opt_eps float 1e-10 Positive Adam-family epsilon.
opt_amsgrad bool False Adam-family AMSGrad option.
opt_sgd_momentum float 0.0 Non-negative SGD momentum.
opt_sgd_dampening float 0.0 Non-negative SGD dampening.
opt_sgd_nesterov bool False Requires positive momentum and zero dampening.

Implementation, precision, and memory

Field Declared type Default Meaning / effective behavior
opt_fused_rotation bool True Request native fused rotations when eligible. Not a guarantee that every scope/device uses them.
opt_gradient_checkpointing Optional[bool] None Scope-dependent resolution above; explicit booleans override it.
opt_stage_cudagraph bool True Request eligible CUDA-graph stage execution. Full-model module processing explicitly disables per-linear stage graphs to avoid private-pool memory growth.
opt_amp Optional[bool] None Standalone/module-stage resolver: CUDA AMP by default except on ROCm; no AMP on non-CUDA devices. Grouped processor runners instead derive AMP directly from device.type == "cuda".
opt_best_state_dtype Union[str, torch.dtype] "fp32" Normalized to "fp16", "bf16", or "fp32" for best-state snapshots.
opt_stage_impl str "fast" "fast" or "reference" for the standalone/module optimizer stage runner. Grouped processing has its own runners.
opt_pair_impl str "fast" "fast" or "reference" pair generation.
opt_quantizer_impl str "reference" "fast" or "reference"; also changes internal quantizer symmetry as explained above.
opt_channel_scale_clamp_min float 1e-2 Positive lower bound for learned channel scales.
opt_channel_scale_clamp_max float 1e2 Positive upper bound; must exceed the lower bound.

quant_linear_init_kwargs() returns {"krot": self.krot}. Configuration serialization records the ParoQuant optimization controls in metadata; beta values serialize as a list. PAROQUANT_EXPORT_FORMATS contains only FORMAT.PAROQUANT. Dynamic format overrides are still required to resolve to FORMAT.PAROQUANT.

Optimization internals

Two-stage loop

optimize_paroquant_linear in gptqmodel/quantization/paroquant/optimization.py:

  1. Validates symmetry and group size; flattens/subsamples captured activation rows. Oversized input streams are sampled using rounded, evenly spaced indices.
  2. Makes FP32 optimization weights, activations, and floating-point targets with F.linear.
  3. Uses a prefix for training and a suffix for validation. With insufficient rows, the two can overlap; this is not an automatically disjoint held-out dataset.
  4. Creates pair metadata, zero angles, unit channel scales, and a mask for dummy rotations.
  5. Optimizes only channel scales and angles during the rotation stage.
  6. Initializes GroupLinearQuantizer from the transformed weight.
  7. Optimizes weights and quantizer parameters during fine-tuning, leaving the rotation parameters inactive.
  8. Exports pseudo-weights, packing weights, quantization parameters, rotation metadata, and losses.

Stage runners use smooth-L1 reconstruction loss, cosine learning-rate decay toward base_lr / 20, per-epoch validation, and restoration of the best validation state. Dummy angles are reset after updates. Zero epochs skip updates, not necessarily quantization/export work.

The return contract is ParoQuantOptimizationResult: pseudo_weight, pack_weight, q_scales, q_zeros, pairs, theta, channel_scales, train_loss, val_loss, and used_identity. Missing activation rows can produce an identity-rotation quantized result, rather than an unquantized layer; processor-level fallback policy determines whether missing features are accepted.

The exact optimizer entry-point signature is:

def optimize_paroquant_linear(
    *,
    weight: torch.Tensor,
    bias: Optional[torch.Tensor],
    inputs: torch.Tensor,
    bits: int,
    group_size: int,
    sym: bool,
    krot: int,
    pair_ratio: float,
    train_rows: int,
    val_rows: int,
    batch_size: int,
    rotation_epochs: int,
    finetune_epochs: int,
    rotation_lr: float,
    weight_lr: float,
    quantizer_lr: float,
    seed: int,
    optimizer_name: str = "adamw",
    optimizer_weight_decay: float = 0.01,
    optimizer_betas: tuple[float, float] = (0.9, 0.95),
    optimizer_eps: float = 1e-10,
    optimizer_amsgrad: bool = False,
    sgd_momentum: float = 0.0,
    sgd_dampening: float = 0.0,
    sgd_nesterov: bool = False,
    fused_rotation: Optional[bool] = None,
    stage_cudagraph: Optional[bool] = None,
    stage_impl: Literal["fast", "reference"] = "fast",
    pair_impl: Literal["fast", "reference"] = "fast",
    quantizer_impl: Literal["fast", "reference"] = "fast",
    gradient_checkpointing: bool = False,
    best_state_dtype: Optional[str | torch.dtype] = "fp32",
    use_amp: Optional[bool] = None,
    scale_clamp_min: float = PAROQUANT_OPT_SCALE_CLAMP_MIN_DEFAULT,
    scale_clamp_max: float = PAROQUANT_OPT_SCALE_CLAMP_MAX_DEFAULT,
) -> ParoQuantOptimizationResult:

Notice that this low-level function defaults quantizer_impl to "fast", whereas ParoConfig.opt_quantizer_impl defaults to "reference". The processor passes the configuration explicitly.

Processor lifecycle

ParoQuantProcessor in gptqmodel/looper/paroquant_processor.py registers modules, captures their input activations to CPU, tracks subset completion, then quantizes a layer after its capture is complete. It rejects calibration-dataset replacement and requires an initial calibration dataset.

For grouped scopes it preserves layer inputs, kwargs, and clean targets. Replay removes past-key-value caches, disables use_cache, rebuilds position/rotary metadata, and normalizes inference-born tensors so autograd can use them. Compute-block optimization uses layer clones. Eligible dense layer-scope optimization instead operates on the live layer, streams CPU-owned replay shards, snapshots best state on CPU, temporarily uses eager attention, and restores the original modules/dtype/attention configuration afterward. The live-layer eligibility check is restricted to the recognized dense projection names and excludes expert-like modules; other layer groups use the clone path.

Layer scope disables fused rotation in its optimizer wrappers. Grouped runners currently choose AMP from the CUDA device type rather than consulting opt_amp; do not assume all low-level controls affect all scopes identically.

Module seeds are derived with BLAKE2b from the base seed, layer index, and a name key. Module scope uses the full module name; grouped scopes use the terminal archetype. This is stronger isolation than simply giving every layer the same seed.

Packing installs ParoLinear, packs the transformed pseudo-quantized weight, copies pairs, theta, and runtime channel_scales, then calls post_init. Finalization marks the model quantized and sets METHOD.PARO. Optional original-versus-quantized residuals can feed a subsequent EoRA processor.

STE acceleration kernels

Source Role Requirements / behavior
optimization.py Torch round/clamp STE, affine custom-autograd node, group quantizer, optional compiled training quantization Reference/fallback arithmetic and orchestration.
modules/triton/affine_ste.py Fused affine pseudo-quantization Optional Triton; selects CUDA or HIP libdevice; ties-to-even rounding and round-to-nearest division. Defaults to a Torch elementwise backward; optional Triton backward still performs group reductions with Torch.
modules/triton/sym_ste.py Fused symmetric pseudo-quantization Optional Triton; fused forward and Torch backward reductions.
utils/paroquant.py Native fused rotation and custom backward Optional JIT extension; supports input, angle, and channel-scale gradients when dispatched through the autograd wrapper.

GPTQMODEL_PAROQUANT_AFFINE_STE_TRITON and GPTQMODEL_PAROQUANT_SYM_STE_TRITON default to enabled. GPTQMODEL_PAROQUANT_AFFINE_STE_TRITON_BACKWARD defaults to disabled. Availability probes and wrapper eligibility checks determine whether these accelerators actually run; they are distinct from the inference Triton GEMM kernels.

Checkpoint and inference implementation

Stored tensors

For a successfully packed four-bit module, let I = in_features, O = out_features, and G be the normalized group size (I for group_size=-1). The AWQ INT4 packer requires O divisible by 8, even though ParoLinear's class declaration advertises output divisibility by 1; declaring a shape supported does not remove the packer's constraint.

Tensor Meaning / shape
qweight AWQ-layout INT32-packed four-bit transformed weights, (I, O // 8).
qzeros INT32-packed group zero points, (I // G, O // 8).
scales FP16/BF16 group dequantization scales, (I // G, O).
pairs INT16 pair schedule, (krot, in_features).
theta Rotation angles, (krot, in_features // 2).
channel_scales Runtime activation scaling, (1, in_features).
bias Optional floating-point linear bias, (O,).

Identity state has zero angles and unit channel scales. A loader that only knows AWQ weight packing but ignores these extra tensors does not implement ParoQuant inference.

Backend reference

Backend / class Declared support Execution
BACKEND.PAROQUANT_CUDA / ParoLinear 4 bits; sym=True; groups -1, 16, 32, 64, 128; INT32 packing; FP16/BF16 inputs; Lora; class declares all devices/platforms Rotation followed by eligible AWQ CUDA GEMM, CPU AVX512, CPU ATen, or temporary dense fallback. The backend name does not mean every call uses a CUDA kernel.
BACKEND.PAROQUANT_TRITON / ParoQuantTritonLinear CUDA and ROCm device types; Linux/Windows; Triton >=2.0.0, or >=3.4.0 with GIL disabled; inherits four-bit/symmetry/dtype/adapter declarations Autotunes decode-fused, prefill-fused, and temporary-dequantize-plus-matmul plans. Low-level GEMM helpers accept groups 32, 64, 128 or full width; group 16 is inherited at class level but not supported by those helpers. The kernels use only portable Triton ops, which is what the ROCm declaration is based on; AUTO still never picks it (SUPPORTS_FORMATS[FORMAT.PAROQUANT] = 0), so select it explicitly.
BACKEND.MLX / MlxParoLinear MLX/Metal implementation for Apple Silicon MLX quantized matmul after staged or fused Metal rotation; returns the activation dtype. See MLX on Apple Silicon for loader requirements.

ParoLinear.forward flattens leading activation dimensions, rotates inputs, evaluates the quantized projection, applies an optional adapter to the original unrotated input, and restores the output shape. Dense dequantization is temporary; it is not a persistent dense-weight cache.

utils/paroquant.py offers reference Torch rotations and a vectorized CPU rotation. The native extension is built from gptqmodel_ext/paroquant/rotation.cu through the shared JIT loader. Its fused rotation eligibility is narrower than the Python API: group size 128, krot 1 or 8, and eligible CUDA tensors. Unsupported cases fall back to Torch. Building the native extension requires the relevant Torch extension toolchain; merely installing Triton does not build it.

def apply_paroquant_rotation(
    x: torch.Tensor,
    pairs: torch.Tensor,
    theta: torch.Tensor,
    scales: Optional[torch.Tensor] = None,
    group_size: int = 128,
) -> torch.Tensor:

Use apply_paroquant_rotation_autograd for its custom fused-backward dispatch; the reference and vectorized helpers expose the same argument structure. Rotation launch selection has process-local and optional disk caches keyed by runtime/hardware identity. These caches store launch decisions, not expanded weights.

Native runtime controls

Environment variable Default Purpose
GPTQMODEL_PAROQUANT_CACHE_RUNTIME_DTYPE False Retain runtime-typed scales/bias rather than repeatedly converting them.
GPTQMODEL_PAROQUANT_AUTO_CACHE_BF16_RUNTIME_DTYPE True Enable that behavior automatically for BF16 compute.
GPTQMODEL_PAROQUANT_CACHE_ROTATION_DTYPE False Cache typed angle and channel-scale metadata.
GPTQMODEL_PAROQUANT_AUTO_CACHE_BF16_ROTATION_DTYPE True Enable typed rotation metadata caching for BF16.
GPTQMODEL_PAROQUANT_CPU_AVX512 True Permit the eligible BF16 CPU AVX512 path.
GPTQMODEL_PAROQUANT_ROTATE_AUTOTUNE True Autotune the native rotation launch when no valid explicit launch override is supplied.
GPTQMODEL_PAROQUANT_ROTATE_CTA_M Unset Explicit value from 4, 8, 16; invalid values are ignored.
GPTQMODEL_PAROQUANT_ROTATE_ROW_PAD Unset Explicit value 0 or 2; invalid values are ignored.
GPTQMODEL_PAROQUANT_AUTOTUNE_CACHE True Enable disk persistence of rotation launch decisions.
GPTQMODEL_PAROQUANT_BUILD_ROOT ~/.cache/gptqmodel/torch_extensions/paroquant unless the shared root is overridden Override the native extension build root. The shared override GPTQMODEL_TORCH_EXTENSIONS_DIR uses its paroquant subdirectory.
GPTQMODEL_PAROQUANT_FORCE_REBUILD False Request native extension rebuilding through the shared loader.
GPTQMODEL_EXT_VERBOSE False Enable shared JIT-extension build verbosity.

Runtime dtype flags are read at module import; constructor arguments with corresponding names can override the dtype-cache settings. The CPU rotation path also reuses typed rotation metadata directly. Metadata caches are not dequantized weight caches. The disk launch-cache root is GPTQMODEL_CACHE_DIR, otherwise $XDG_CACHE_HOME/gptqmodel, otherwise ~/.cache/gptqmodel; files live under paroquant/rotation_launch/. Use the rotation cache/extension clearing helpers in utils/paroquant.py when deliberately changing tuning or build state within one process.

Triton plans and tuning

The low-level GEMM unpacks INT4 tiles, subtracts zero points, multiplies by group scales, and accumulates the dot product. Decode uses blocks (M, N, K)=(4, 128, 32), four warps, and two stages; prefill uses (32, 128, 32), eight warps, and four stages. GPTQMODEL_FP32_ACCUM defaults to enabled.

Environment variable Default Meaning
GPTQMODEL_PAROQUANT_TRITON_AUTOTUNE "1" Any value other than "0" enables plan tuning.
GPTQMODEL_PAROQUANT_TRITON_AUTOTUNE_WARMUP 2 Warmups, clamped to at least zero.
GPTQMODEL_PAROQUANT_TRITON_AUTOTUNE_ITERS 4 Timed iterations, at least one.
GPTQMODEL_PAROQUANT_TRITON_AUTOTUNE_MARGIN 0.05 Non-negative tuning margin.
GPTQMODEL_PAROQUANT_TRITON_DECODE_MAX_ROWS 8 Decode-row threshold, at least one.

These class settings are read at import time. Set them before importing the runtime. A module's clear_autotune() clears its plan decisions. The plan cache is per module and keyed only by "decode" or "prefill", not by the exact activation shape; the first tuned workload of each kind determines later calls until the cache is cleared.

ROCm and Apple Silicon distinctions

Torch uses the cuda device type for ROCm too. The standalone optimizer explicitly disables stage CUDA graphs on ROCm and defaults AMP off there. This is separate from the grouped runners' device-type-based AMP behavior. Triton training kernels import HIP libdevice on HIP builds; the inference class validates torch.cuda.is_available() rather than establishing compatibility with every AMD architecture. Do not infer a hardware-support guarantee from a shared device-type name.

The MLX optimizer is available only on Darwin/arm64 with mlx.core importable. ParoQuantProcessor selects it for nonempty CPU/MPS module-scope calibration when opt_stage_impl="fast" and gradient checkpointing is off. It implements the two stages with MLX autodiff and bridges tensors through DLPack. Grouped scopes are not this native MLX optimization path.

MLX source Responsibility
quantization/mlx_paroquant_optimize.py Availability, differentiable rotations/STE, AdamW/Adam/SGD, stage loop, Torch result bridge.
quantization/mlx_paroquant_rotation.py Forward/inverse Metal rotation, validated local pair permutations; standalone primitive returns FP32. Groups 16, 32, 64, 128.
quantization/mlx_paroquant_quant.py Metal transformed-weight pseudo-quantization; primitive supports bits 2..8, distinct from the four-bit export runtime.
quantization/mlx_paroquant.py Native symmetric four-bit AWQ-layout packing; validates finite positive scales, finite weights, code range, and output width divisible by 32.
nn_modules/qlinear/mlx_paro.py Runtime staged/fused Metal activation rotations and quantized matmul. Uses FP32 intermediates and restores the input dtype.

The MLX runtime uses fused rotations for BF16 or multiple stages when the group is at most 128; eligible single-stage FP16 execution may use the staged kernel.

Benchmark and development utilities

gptqmodel/utils/paroquant_benchmark.py provides dense baselines, prefix-layer/single-module/selected-module dynamic masks, quantize-save-reload evaluation, module/region timing tables, and CUDA-versus-Triton-versus-dense kernel comparisons. Kernel measurements capture real module inputs and report maximum/mean absolute differences as well as latency. Evaluation helpers are imported from tests/eval.py, so this is checkout tooling, not a self-contained installed-package benchmark.

Its default model/calibration locations are local /monster/data/model/... paths. Adapt them to your installation before running it. The utility's default row counts and evaluation suite are benchmark choices, not additional public quantizer defaults. See Evaluation for model-quality measurement.

GSQ refinement

ParoConfig.gsq adds a Gumbel-Softmax refinement pass to ParoQuant's module optimization (gsq_paro.refine_paro_export). It requires bits 4, opt_scope "module", and opt_train_on_noisy_inputs False. See GSQ.

Pitfalls and edge cases

  • Keep sym=True and use the ParoQuant format. Do not relabel the checkpoint as AWQ.
  • The optimizer's generic pseudo-quantizer can handle more bit widths than the current four-bit inference/export backend. Helper acceptance is not end-to-end format support.
  • Valid group sizes must also divide the actual module width and satisfy the selected kernel's restrictions. Automatic padding is not advertised by the ParoQuant runtime classes.
  • Calibration rows are required for meaningful learned rotations. Identity fallback is not successful rotation optimization; disable fallback when missing activations should be an error.
  • Module scope reconstructs linears; grouped scopes reconstruct a whole-layer output. Their losses and memory costs are not interchangeable measurements.
  • Validation prefixes/suffixes can overlap when data is insufficient. Supply enough calibration rows if a disjoint holdout matters.
  • More stages increase rotation metadata and work; they do not expand the fused native kernel's supported krot set.
  • CUDA graphs and AMP are scope- and platform-dependent, even when enabled in configuration.
  • Changing opt_quantizer_impl changes quantization arithmetic, not just speed.
  • Keep rotation buffers with the checkpoint and preserve krot when constructing runtime modules.

Source files

Paths below are relative to the repository:

  • gptqmodel/quantization/config.py — configuration, method/format values, metadata.
  • gptqmodel/quantization/paroquant/optimization.py — optimization, STE, pairing, result contract.
  • gptqmodel/quantization/paroquant/modules/triton/affine_ste.py — affine STE acceleration.
  • gptqmodel/quantization/paroquant/modules/triton/sym_ste.py — symmetric STE acceleration.
  • gptqmodel/quantization/paroquant/modules/triton/gemm.py — inference dequantization/GEMM.
  • gptqmodel/looper/paroquant_processor.py — capture, scope orchestration, packing.
  • gptqmodel/nn_modules/qlinear/paroquant.py — shared runtime.
  • gptqmodel/nn_modules/qlinear/paroquant_triton.py — Triton plan dispatch.
  • gptqmodel/nn_modules/qlinear/torch_awq.py, gptqmodel/quantization/awq/utils/packing_utils.py — inherited packing and four-bit layout constraints.
  • gptqmodel/utils/paroquant.py — rotations, JIT/autotuning, autograd.
  • gptqmodel/utils/cpp.py — shared JIT build-root, rebuild, and verbosity policy.
  • gptqmodel/utils/paroquant_benchmark.py — development benchmarks.
  • gptqmodel_ext/paroquant/rotation.cu — native rotation extension.
  • gptqmodel/quantization/mlx_paroquant.py — MLX packing.
  • gptqmodel/quantization/mlx_paroquant_optimize.py — MLX optimizer.
  • gptqmodel/quantization/mlx_paroquant_quant.py — MLX pseudo-quantization.
  • gptqmodel/quantization/mlx_paroquant_rotation.py — MLX weight rotation.
  • gptqmodel/nn_modules/qlinear/mlx_paro.py — MLX runtime.
  • gptqmodel/models/base.py — quantization processor selection.
  • docs/paroquant.md — scope-specific checkpointing notes.

See also

Documentation home · Quickstart · Troubleshooting · Repository · PyPI

Source baseline: GPT-QModel 7.7.0 / c197691. Maintained by ModelCloud.ai.

Home

Getting started

Quantization methods

Complementary configs

Using a model

Platforms

Internals

Clone this wiki locally