PyTorch quantization

PyTorch quantization#

Quark Quantization API for PyTorch.

class quark.torch.quantization.api.ModelQuantizer(config: QConfig, multi_device: bool = False)[source]#

Provides an API for quantizing deep learning models using PyTorch.

This class handles the configuration and processing of the model for quantization based on user-defined parameters. It is essential to ensure that the ‘config’ provided has all necessary quantization parameters defined. This class assumes that the model is compatible with the quantization settings specified in ‘config’.

Parameters:: config (QConfig) – The model quantization configuration.

quantize_model(model: Module, dataloader: DataLoader[Tensor] | DataLoader[list[dict[str, Tensor]]] | DataLoader[dict[str, Tensor]] | DataLoader[list[BatchFeature]] | None = None) → Module[source]#

Quantizes the given PyTorch model to optimize its performance and reduce its size.

The dataloader is used to provide data necessary for calibration during the quantization process. Depending on the type of data provided (either tensors directly or structured as lists or dictionaries of tensors), the function will adapt the quantization approach accordingly.

It is important that the model and dataloader are compatible in terms of the data they expect and produce. Misalignment in data handling between the model and the dataloader can lead to errors during the quantization process.

Parameters:

model (torch.nn.Module) – The PyTorch model to be quantized. This model should be already trained and ready for quantization.
dataloader (Optional[Union[DataLoader[torch.Tensor], DataLoader[List[Dict[str, torch.Tensor]]], DataLoader[Dict[str, torch.Tensor]], DataLoader[List[BatchFeature]]]]) – The torch.utils.data.DataLoader providing data that the quantization process will use for calibration. This can be a simple DataLoader returning tensors, or a more complex structure returning either a list of dictionaries or a dictionary of tensors.

Returns:

The quantized version of the input model. This model is now optimized for inference with reduced size and potentially improved performance on targeted devices.

Return type:

torch.nn.Module

Example:

# Model & Data preparation
from torch.utils.data import DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer

from quark.torch.quantization.config.config import QConfig
from quark.torch.quantization.config.type import Dtype, ScaleType, RoundType, QSchemeType
from quark.torch.quantization.observer.observer import PerGroupMinMaxObserver

from quark.torch import ModelQuantizer

model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m", torch_dtype="auto")
model.eval()
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")

quant_spec = QTensorConfig(
    dtype=Dtype.uint4,
    observer_cls=PerGroupMinMaxObserver,
    symmetric=False,
    scale_type=ScaleType.float,
    round_method=RoundType.half_even,
    qscheme=QSchemeType.per_group,
    ch_axis=1,
    is_dynamic=False,
    group_size=128
)
quant_config = QConfig(global_quant_config=QLayerConfig(weight=quant_spec))

text = "Hello, how are you?"
tokenized_outputs = tokenizer(text, return_tensors="pt")
calib_dataloader = DataLoader(tokenized_outputs['input_ids'])

quantizer = ModelQuantizer(quant_config)
quant_model = quantizer.quantize(model, calib_dataloader)

direct_quantize_checkpoint(pretrained_model_path: str, save_path: str, keep_excluded_layers_as_original_model_state: bool = False, *legacy_device_args: str | device, weight_converters: list[Any] | None = None, device: str | device | None = None, presharded_weights: dict[str, int] | None = None) → None[source]#

Quantize model weights by processing each safetensors file independently (file-to-file mode).

This method provides memory-efficient weight-only quantization by processing safetensors files one at a time, without loading the full model into GPU memory. This is particularly useful for quantizing very large models that exceed available GPU memory.

The quantized shards and all configuration files (config.json, model.safetensors.index.json, tokenizer files, etc.) are written to save_path.

Parameters:

pretrained_model_path (str) – Path to the pretrained model directory containing safetensors files.
save_path (str) – Directory path to save the quantized safetensors files.
keep_excluded_layers_as_original_model_state (bool) – If True, tensors already quantized in the source checkpoint but excluded from Quark quantization keep their original model-state format in the export. Defaults to False.
weight_converters (list | None) – Optional list of WeightConverter instances to transform tensors after precision recovery and before quantization. For example, splitting fused gate_up_proj into separate gate_proj and up_proj. Defaults to None.
device (str | torch.device) – Device for tensor operations (e.g., "cuda", "cuda:0", "cpu"). Defaults to "cuda". Legacy positional callers may still pass device after keep_excluded_layers_as_original_model_state.

Example:

from quark.torch.quantization.config.config import QConfig, QLayerConfig, OCP_MXFP4Spec
from quark.torch.quantization.weight_convert import Chunk, WeightConverter

from quark.torch import ModelQuantizer

weight_spec = OCP_MXFP4Spec(ch_axis=-1, is_dynamic=False).to_quantization_spec()
quant_config = QConfig(global_quant_config=QLayerConfig(weight=weight_spec))

weight_converters = [
    WeightConverter(
        "gate_up_proj.weight",
        ["gate_proj.weight", "up_proj.weight"],
        operations=[Chunk(dim=0)],
    ),
]

quantizer = ModelQuantizer(quant_config)
quantizer.direct_quantize_checkpoint(
    pretrained_model_path="/path/to/model",
    save_path="/path/to/output",
    weight_converters=weight_converters,
)

static freeze(model: Module | GraphModule, quantize: bool | None = None, runtime_options: RuntimeOptions | None = None) → Module | GraphModule[source]#

Freezes the quantized model by replacing FakeQuantize modules with FrozenFakeQuantize modules.

In order to be able to compile a quantized model through torch.compile, this method needs to be applied.

Parameters:

model (torch.nn.Module) – The neural network model containing quantized layers.
quantize (Optional[bool]) – Whether to effectively quantize weights, moving away from soft weights that are quantized on the fly to e.g. QuantLinear.weight actually holding the fake quantized weights. This can be disabled e.g. if we would like simply to move to use FrozenFakeQuantize from a model using QuantLinear that is already holding the fake quantized weights in high-precision. Defaults to True for PyTorch eager models, and False for FX graph models.
runtime_options (Optional[RuntimeOptions]) – Runtime toggles controlling native inference conversion (native linear mode selection, preshuffle, in-place strategy). If provided, native inference conversion is enabled after freeze.

Returns:

The modified model with FakeQuantize modules replaced by FrozenFakeQuantize modules.

Return type:

torch.nn.Module

quark.torch.quantization.api.load_params(model: Module | None = None, json_path: str = '', safetensors_path: str = '', pth_path: str = '', quant_mode: QuantizationMode = QuantizationMode.eager_mode, compressed: bool = False, reorder: bool = True) → Module[source]#

Instantiates a quantized model from saved model files, which is generated from the quark.torch.export.api.save_params() function.

Parameters:

model (torch.nn.Module) – The original Pytorch model.
json_path (str) – The path of the saved json file. Only available for eager mode quantization.
safetensors_path (str) – The path of the saved safetensors file. Only available for eager mode quantization.
pth_path (str) – The path of the saved .pth file. Only available for fx_graph mode quantization.
quant_mode (QuantizationMode) – The quantization mode. The choice includes "QuantizationMode.eager_mode" and "QuantizationMode.fx_graph_mode". Default is "QuantizationMode.eager_mode".
compressed (bool) – Whether the quantized model to load is stored using its compressed data type, or in a “fake quantized” format (QDQ).
reorder (bool) – Reorder.

Returns:

The reloaded quantized version of the input model.

Return type:

torch.nn.Module

Examples:

# eager mode:
from quark.torch import load_params
model = load_params(model, json_path=json_path, safetensors_path=safetensors_path)

# fx_graph mode:
from quark.torch.quantization.api import load_params
model = load_params(pth_path=model_file_path, quant_mode=QuantizationMode.fx_graph_mode)

Note:: This function does not support dynamic quantization for now.

quark.torch.quantization.api.enable_native_inference(model: Module, *, runtime_options: RuntimeOptions | None = None) → int[source]#

Enable native inference mode for quantized linear layers in the model.

Replaces eligible QuantLinear / QParamsLinear modules with Aiter-backed native inference layers that use optimized AMD GEMM kernels for native inference. Export compatibility is preserved through shared state_dict serialization helpers.

This can be called directly on a quantized model: QuantLinear -> NativeInferenceLinear conversion is handled internally.

When preshuffle is enabled (via runtime_options.use_preshuffle), the weight is shuffled in-place (~1× weight memory); state_dict() temporarily un-shuffles for export.

To revert back to the export-format module (QParamsLinear with the scaled_mm / dequant fallback forward path), call disable_native_inference() directly — there is no boolean toggle.

Parameters:

model (torch.nn.Module) – The quantized model containing QuantLinear / QParamsLinear layers.
runtime_options (RuntimeOptions) – Runtime conversion options, including native linear selection and preshuffle behavior.

Returns:

Number of layers converted.

Return type:

int

Raises:

ImportError – If Aiter is not installed.

quark.torch.quantization.api.disable_native_inference(model: Module) → int[source]#

Convert native inference layers back to base QParamsLinear.

The base QParamsLinear.forward() uses scaled_mm (FP8 per-tensor) or dequant + F.linear as a fallback, so no Aiter dependency is needed after disabling.

Parameters:: model (torch.nn.Module) – The model containing native inference layers.
Returns:: Number of layers converted back to base QParamsLinear.
Return type:: int

class quark.torch.quantization.api.RuntimeOptions(native_linear_mode: str = 'auto', use_preshuffle: bool = False, svdquant_overlap_streams: bool = False)[source]#: Runtime toggles for native inference conversion.

PyTorch quantization

Contents

PyTorch quantization#