provides high-performance dtypes and optimizations based on quantization and sparsity for inference and training PyTorch models. It is supported for any model in any modality, as long as it supports loading with
and contains torch.nn.Linear layers.
Make sure Pytorch 2.5+ and torchao are installed with the command below.
uv pip install -U torch torchaoEach quantization dtype is available as a separate instance of a
class. This provides more flexible configuration options by exposing more available arguments.
Pass the AOBaseConfig of a quantization dtype, like
to [TorchAoConfig] in [~ModelMixin.from_pretrained].
importtorchfromdiffusersimportDiffusionPipeline, PipelineQuantizationConfig, TorchAoConfigfromtorchao.quantizationimportInt8WeightOnlyConfigpipeline_quant_config=PipelineQuantizationConfig( quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128, version=2))} ) pipeline=DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", quantization_config=pipeline_quant_config, dtype=torch.bfloat16, device_map="cuda"# or "mps", "xpu", "cpu" )device_map="cuda" quantizes each layer on the GPU while it loads. This is fast, but it temporarily requires additional GPU memory for the original and quantized weights. If the model already uses most of your GPU memory, loading can fail with an out-of-memory error. In that case, drop device_map and move the pipeline to the GPU after loading:
pipeline=DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16, ) pipeline.to("cuda") # or "mps", "xpu", "cpu"Without device_map, Diffusers quantizes the layers on the CPU. This is slower, but avoids the temporary GPU-memory spike during quantization. To reduce GPU memory usage further, use [~DiffusionPipeline.enable_model_cpu_offload] instead. You can also quantize additional components, such as the text encoder.
torch.compile
torchao supports
which can speed up inference with one line of code.
importtorchfromdiffusersimportDiffusionPipeline, PipelineQuantizationConfig, TorchAoConfigfromtorchao.quantizationimportInt4WeightOnlyConfigpipeline_quant_config=PipelineQuantizationConfig( quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))} ) pipeline=DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", quantization_config=pipeline_quant_config, dtype=torch.bfloat16, device_map="cuda"# or "mps", "xpu", "cpu" ) pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True)Refer to this
for inference speed and memory usage benchmarks with Flux and CogVideoX. More benchmarks on various hardware are also available in the torchao
.
Tip
The FP8 post-training quantization schemes in torchao are effective for GPUs with compute capability of at least 8.9 (RTX-4090, Hopper, etc.). FP8 often provides the best speed, memory, and quality trade-off when generating images and videos. We recommend combining FP8 and torch.compile if your GPU is compatible.
Supported quantization types
torchao supports weight-only quantization and weight and dynamic-activation quantization for int8, float3-float8, and uint1-uint7.
Weight-only quantization stores the model weights in a specific low-bit data type but performs computation with a higher-precision data type, like bfloat16. This lowers the memory requirements from model weights but retains the memory peaks for activation computation.
Dynamic activation quantization stores the model weights in a low-bit dtype, while also quantizing the activations on-the-fly to save additional memory. This lowers the memory requirements from model weights, while also lowering the memory overhead from activation computations. However, this may come at a quality tradeoff at times, so it is recommended to test different models thoroughly.
Refer to the
official torchao documentation
for a better understanding of the available quantization methods. An exhaustive list of configuration options are available
.
Some example popular quantization configurations are as follows:
CategoryConfiguration ClassesInteger quantization
,
,
Int8DynamicActivationInt8WeightConfig
Floating point 8-bit quantization
,
Float8DynamicActivationFloat8WeightConfig
Unsigned integer quantization
Serializing and Deserializing quantized models
To serialize a quantized model in a given dtype, first load the model with the desired quantization dtype and then save it using the [~ModelMixin.save_pretrained] method.
importtorchfromdiffusersimportAutoModel, TorchAoConfigfromtorchao.quantizationimportInt8WeightOnlyConfigquantization_config=TorchAoConfig(Int8WeightOnlyConfig()) transformer=AutoModel.from_pretrained( "black-forest-labs/Flux.1-Dev", subfolder="transformer", quantization_config=quantization_config, dtype=torch.bfloat16, ) transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)To load a serialized quantized model, use the [~ModelMixin.from_pretrained] method.
importtorchfromdiffusersimportFluxPipeline, AutoModeltransformer=AutoModel.from_pretrained("/path/to/flux_int8wo", dtype=torch.bfloat16, use_safetensors=False) pipe=FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, dtype=torch.bfloat16) pipe.to("cuda") # or "mps", "xpu", "cpu"prompt="A cat holding a sign that says hello world"image=pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0] image.save("output.png")If you are using torch<=2.6.0, some quantization methods, such as uint4 weight-only, cannot be loaded directly and may result in an UnpicklingError when trying to load the models, but work as expected when saving them. In order to work around this, one can load the state dict manually into the model. Note, however, that this requires using weights_only=False in torch.load, so it should be run only if the weights were obtained from a trustable source.
importtorchfromaccelerateimportinit_empty_weightsfromdiffusersimportFluxPipeline, AutoModel, TorchAoConfigfromtorchao.quantizationimportIntxWeightOnlyConfig# Serialize the modeltransformer=AutoModel.from_pretrained( "black-forest-labs/Flux.1-Dev", subfolder="transformer", quantization_config=TorchAoConfig(IntxWeightOnlyConfig(dtype=torch.uint4)), dtype=torch.bfloat16, ) transformer.save_pretrained("/path/to/flux_uint4wo", safe_serialization=False, max_shard_size="50GB") # ...# Load the modelstate_dict=torch.load("/path/to/flux_uint4wo/diffusion_pytorch_model.bin", weights_only=False, map_location="cpu") withinit_empty_weights(): transformer=AutoModel.from_config("/path/to/flux_uint4wo/config.json") transformer.load_state_dict(state_dict, strict=True, assign=True)Tip
The [AutoModel] API is supported for PyTorch >= 2.6 as shown in the examples below.
Resources