diffusers/pr_10312/en/quantization/torchao.html

<meta charset="utf-8" /><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;torchao&quot;,&quot;local&quot;:&quot;torchao&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Supported quantization types&quot;,&quot;local&quot;:&quot;supported-quantization-types&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Resources&quot;,&quot;local&quot;:&quot;resources&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2}],&quot;depth&quot;:1}">
		<link href="/docs/diffusers/pr_10312/en/_app/immutable/assets/0.e3b0c442.css" rel="modulepreload">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/entry/start.203b6290.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/scheduler.8c3d61f6.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/singletons.8b179f45.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/index.0997d446.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/paths.3fd58d56.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/entry/app.423ea23f.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/index.da70eac4.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/nodes/0.e544eae6.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/each.e59479a4.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/nodes/205.e2dc2a46.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/CodeBlock.00a903b3.js">
		<link rel="modulepreload" href="/docs/diffusers/pr_10312/en/_app/immutable/chunks/EditOnGithub.1e64e623.js"><!-- HEAD_svelte-u9bgzb_START --><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;torchao&quot;,&quot;local&quot;:&quot;torchao&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;Supported quantization types&quot;,&quot;local&quot;:&quot;supported-quantization-types&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Resources&quot;,&quot;local&quot;:&quot;resources&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2}],&quot;depth&quot;:1}"><!-- HEAD_svelte-u9bgzb_END -->      <p></p>   <h1 class="relative group"><a id="torchao" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#torchao"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>torchao</span></h1> <p data-svelte-h="svelte-3zt1h"><a href="https://github.com/pytorch/ao" rel="nofollow">TorchAO</a> is an architecture optimization library for PyTorch. It provides high-performance dtypes, optimization techniques, and kernels for inference and training, featuring composability with native PyTorch features like <a href="https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html" rel="nofollow">torch.compile</a>, FullyShardedDataParallel (FSDP), and more.</p> <p data-svelte-h="svelte-xmmcvn">Before you begin, make sure you have Pytorch 2.5+ and TorchAO installed.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5   text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START -->pip install -U torch torchao<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-15rn8tp">Quantize a model by passing <a href="/docs/diffusers/pr_10312/en/api/quantization#diffusers.TorchAoConfig">TorchAoConfig</a> to <a href="/docs/diffusers/pr_10312/en/api/models/overview#diffusers.ModelMixin.from_pretrained">from_pretrained()</a> (you can also load pre-quantized models). This works for any model in any modality, as long as it supports loading with <a href="https://hf.co/docs/accelerate/index" rel="nofollow">Accelerate</a> and contains <code>torch.nn.Linear</code> layers.</p> <p data-svelte-h="svelte-1d14wga">The example below only quantizes the weights to int8.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5   text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> FluxPipeline, FluxTransformer2DModel, TorchAoConfig

model_id = <span class="hljs-string">&quot;black-forest-labs/Flux.1-Dev&quot;</span>
dtype = torch.bfloat16

quantization_config = TorchAoConfig(<span class="hljs-string">&quot;int8wo&quot;</span>)
transformer = FluxTransformer2DModel.from_pretrained(
    model_id,
    subfolder=<span class="hljs-string">&quot;transformer&quot;</span>,
    quantization_config=quantization_config,
    torch_dtype=dtype,
)
pipe = FluxPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    torch_dtype=dtype,
)
pipe.to(<span class="hljs-string">&quot;cuda&quot;</span>)

prompt = <span class="hljs-string">&quot;A cat holding a sign that says hello world&quot;</span>
image = pipe(prompt, num_inference_steps=<span class="hljs-number">28</span>, guidance_scale=<span class="hljs-number">0.0</span>).images[<span class="hljs-number">0</span>]
image.save(<span class="hljs-string">&quot;output.png&quot;</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1mb3hq5">TorchAO is fully compatible with <a href="./optimization/torch2.0#torchcompile">torch.compile</a>, setting it apart from other quantization methods. This makes it easy to speed up inference with just one line of code.</p> <div class="code-block relative"><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5   text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-comment"># In the above code, add the following after initializing the transformer</span>
transformer = torch.<span class="hljs-built_in">compile</span>(transformer, mode=<span class="hljs-string">&quot;max-autotune&quot;</span>, fullgraph=<span class="hljs-literal">True</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-exfnd4">For speed and memory benchmarks on Flux and CogVideoX, please refer to the table <a href="https://github.com/huggingface/diffusers/pull/10009#issue-2688781450" rel="nofollow">here</a>. You can also find some torchao <a href="https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks" rel="nofollow">benchmarks</a> numbers for various hardware.</p> <p data-svelte-h="svelte-1fx3nq1">torchao also supports an automatic quantization API through <a href="https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md#autoquantization" rel="nofollow">autoquant</a>. Autoquantization determines the best quantization strategy applicable to a model by comparing the performance of each technique on chosen input types and shapes. Currently, this can be used directly on the underlying modeling components. Diffusers will also expose an autoquant configuration option in the future.</p> <p data-svelte-h="svelte-1eoai4f">The <code>TorchAoConfig</code> class accepts three parameters:</p> <ul data-svelte-h="svelte-1xogsn4"><li><code>quant_type</code>: A string value mentioning one of the quantization types below.</li> <li><code>modules_to_not_convert</code>: A list of module full/partial module names for which quantization should not be performed. For example, to not perform any quantization of the <a href="/docs/diffusers/pr_10312/en/api/models/flux_transformer#diffusers.FluxTransformer2DModel">FluxTransformer2DModel</a>’s first block, one would specify: <code>modules_to_not_convert=[&quot;single_transformer_blocks.0&quot;]</code>.</li> <li><code>kwargs</code>: A dict of keyword arguments to pass to the underlying quantization method which will be invoked based on <code>quant_type</code>.</li></ul>  <h2 class="relative group"><a id="supported-quantization-types" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#supported-quantization-types"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Supported quantization types</span></h2> <p data-svelte-h="svelte-1dy3rwb">torchao supports weight-only quantization and weight and dynamic-activation quantization for int8, float3-float8, and uint1-uint7.</p> <p data-svelte-h="svelte-17x1tdo">Weight-only quantization stores the model weights in a specific low-bit data type but performs computation with a higher-precision data type, like <code>bfloat16</code>. This lowers the memory requirements from model weights but retains the memory peaks for activation computation.</p> <p data-svelte-h="svelte-1nbmql6">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.</p> <p data-svelte-h="svelte-4xexxq">The quantization methods supported are as follows:</p> <table data-svelte-h="svelte-18shdov"><thead><tr><th><strong>Category</strong></th> <th><strong>Full Function Names</strong></th> <th><strong>Shorthands</strong></th></tr></thead> <tbody><tr><td><strong>Integer quantization</strong></td> <td><code>int4_weight_only</code>, <code>int8_dynamic_activation_int4_weight</code>, <code>int8_weight_only</code>, <code>int8_dynamic_activation_int8_weight</code></td> <td><code>int4wo</code>, <code>int4dq</code>, <code>int8wo</code>, <code>int8dq</code></td></tr> <tr><td><strong>Floating point 8-bit quantization</strong></td> <td><code>float8_weight_only</code>, <code>float8_dynamic_activation_float8_weight</code>, <code>float8_static_activation_float8_weight</code></td> <td><code>float8wo</code>, <code>float8wo_e5m2</code>, <code>float8wo_e4m3</code>, <code>float8dq</code>, <code>float8dq_e4m3</code>, <code>float8_e4m3_tensor</code>, <code>float8_e4m3_row</code></td></tr> <tr><td><strong>Floating point X-bit quantization</strong></td> <td><code>fpx_weight_only</code></td> <td><code>fpX_eAwB</code> where <code>X</code> is the number of bits (1-7), <code>A</code> is exponent bits, and <code>B</code> is mantissa bits. Constraint: <code>X == A + B + 1</code></td></tr> <tr><td><strong>Unsigned Integer quantization</strong></td> <td><code>uintx_weight_only</code></td> <td><code>uint1wo</code>, <code>uint2wo</code>, <code>uint3wo</code>, <code>uint4wo</code>, <code>uint5wo</code>, <code>uint6wo</code>, <code>uint7wo</code></td></tr></tbody></table> <p data-svelte-h="svelte-1r64eqr">Some quantization methods are aliases (for example, <code>int8wo</code> is the commonly used shorthand for <code>int8_weight_only</code>). This allows using the quantization methods described in the torchao docs as-is, while also making it convenient to remember their shorthand notations.</p> <p data-svelte-h="svelte-9cn9y3">Refer to the official torchao documentation for a better understanding of the available quantization methods and the exhaustive list of configuration options available.</p>  <h2 class="relative group"><a id="resources" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#resources"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Resources</span></h2> <ul data-svelte-h="svelte-96lbpr"><li><a href="https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md" rel="nofollow">TorchAO Quantization API</a></li> <li><a href="https://github.com/sayakpaul/diffusers-torchao" rel="nofollow">Diffusers-TorchAO examples</a></li></ul> <a class="!text-gray-400 !no-underline text-sm flex items-center not-prose mt-4" href="https://github.com/huggingface/diffusers/blob/main/docs/source/en/quantization/torchao.md" target="_blank"><span data-svelte-h="svelte-1kd6by1">&lt;</span> <span data-svelte-h="svelte-x0xyl0">&gt;</span> <span data-svelte-h="svelte-1dajgef"><span class="underline ml-1.5">Update</span> on GitHub</span></a>  <p></p> 
			
			<script>
				{
					__sveltekit_zvhs44 = {
						assets: "/docs/diffusers/pr_10312/en",
						base: "/docs/diffusers/pr_10312/en",
						env: {}
					};

					const element = document.currentScript.parentElement;

					const data = [null,null];

					Promise.all([
						import("/docs/diffusers/pr_10312/en/_app/immutable/entry/start.203b6290.js"),
						import("/docs/diffusers/pr_10312/en/_app/immutable/entry/app.423ea23f.js")
					]).then(([kit, app]) => {
						kit.start(app, element, {
							node_ids: [0, 205],
							data,
							form: null,
							error: null
						});
					});
				}
			</script>