Quantization libraries have a reputation for being bolt-on afterthoughts: a separate graph mode, a foreign config format, kernels that break the moment you wrap your model in torch.compile. TorchAO, launched by the PyTorch team in September 2024 and versioned 0.19.0 in this tree, takes the opposite approach. It is pure native PyTorch, BSD 3-Clause licensed, and it works out of the box with torch.compile and FSDP2 across most HuggingFace models. The numbers on its README are the reason it ended up as the default quantization backend for HuggingFace Transformers, vLLM, ExecuTorch, and diffusers: Llama-3.1-70B pre-training runs 1.5x faster with float8 training, QAT recovers 67 percent of quantized accuracy degradation on Gemma3-4B, and an int4-quantized Llama-3-8B serves 1.89x faster inference with 58 percent less memory. As always in this series, what follows is an educational tour of published source code.

The design idea that holds the whole repository together is a two-layer split. The top layer is a small grammar of configs: dataclasses like Int4WeightOnlyConfig or Float8DynamicActivationFloat8WeightConfig that describe what you want, not how to do it. The bottom layer is a family of torch.Tensor subclasses, each of which owns its own quantize, dequantize, and matmul dispatch. Between them sits a registry that maps every config type to a function that walks the model and swaps linear weights in place. Everything else, from quantization-aware training to low-bit optimizers to experimental MoE kernels, plugs into that same registry pattern. Reading the overview from left to right traces one linear layer from a config object to a packed four-bit tensor running custom kernels.

Overview of TorchAO: quantize_ consumes AOBaseConfig subclasses with granularity choices, dispatches through the transform registry to tensor subclasses for int4, NF4, and float8, while QAT and float8 training flows and low-bit optimizers feed native C++ and CUDA ops.

Reading the overview from left to right:

• quantize_ is the public entry point at line 299 of an 1817-line quant_api.py that also hosts every inference config dataclass.
• AOBaseConfig at line 27 of a 329-line core module is the abstract parent all workflow configs extend, defining layout and filter hooks.
• The transform registry is a 52-line module holding the handler dict at line 13, with register_quantize_module_handler at 19 letting any package plug in.
• TorchAOBaseTensor at line 721 of 1272-line utils.py is the subclass contract every quantized tensor implements.
• Int4Tensor at line 57 of 627 lines stores packed four-bit weights, with packing variants like tile_packed_to_4d in a 391-line sibling file.
• Float8Tensor at line 84 of 1126 lines carries float8 data plus scales for dynamic float8 inference.
• QATConfig at line 43 of 481 lines drives quantization-aware training through a prepare step and a convert step.
• The low-bit optimizers live in a 430-line file: AdamW8bit at 309, AdamW4bit at 342, AdamWFp8 at 375, all sharing _AdamBase at 19.

Why You Need This

If you serve large language models, memory bandwidth is your tax bracket. A 70-billion-parameter model in bfloat16 needs about 140 gigabytes just for weights, which means multiple accelerators before you serve a single token. TorchAO’s int4 weight-only path compresses those weights to roughly a quarter of the size while keeping the compute in higher precision, and the measured results are concrete: 1.73x speedup with 65 percent less memory for Gemma3-12b-it on an H100, 1.5 to 1.6x on gemma-3-27b-it with float8 dynamic quantization, and 1.54x and 1.27x on Flux.1-Dev and CogVideoX-5b image and video generation respectively. On the mobile side, a Qwen3-4B model quantized to int8 activations with int4 weights runs at 14.8 tokens per second in 3379 megabytes on an iPhone 15 Pro through ExecuTorch.

If you train models, the float8 story matters even more. The float8 data type halves the bytes per element of bfloat16, but naive casting wrecks convergence because float8 has almost no mantissa bits near zero. TorchAO’s training stack handles amax history, per-tensor or per-row scaling, and delayed scaling so that Llama-scale pre-training runs at 1.43 to 1.51x on 70B and 405B models with FSDP2, and rowwise float8 yielded 1.34 to 1.43x at 2k H100 GPU scale with matching loss curves. The newest prototype pushes the same idea to mixture-of-experts: MXFP8 grouped matmuls measured roughly 1.45x on the MoE layer of Llama4 Scout and about 1.25x on DeepSeekV3 671B, with a 1.28x end-to-end result on a Crusoe B200 cluster.

If you do both, the QAT loop is the piece most libraries get wrong. Post-training quantization is cheap but costs accuracy; quantization-aware training recovers it, but many implementations are fragile under torch.compile. TorchAO’s QAT recovered up to 96 percent of the accuracy degradation on Llama-3-8B, 67 percent on Gemma3-4B, and the approach has since been adopted by Unsloth and Axolotl for fine-tuning workflows. Because everything composes with torch.compile and FSDP2, you can float8-train, QAT-finetune, int4-quantize, and serve through vLLM without leaving the PyTorch ecosystem. The January 2025 integration with GemLite and SGLang measured 1.1 to 2x faster inference across batch sizes and tensor parallel widths.

Detail view of TorchAO: config dataclasses extend AOBaseConfig with granularity, the engine walks modules through the handler registry, workflow folders build Int4, Int8, NF4, and Float8 tensors on quant_primitives, the float8 training stack swaps in Float8Linear, and optimizers plus custom ops round out the runtime.

The engine mechanics are worth reading because they are small. quantize_ at line 299 accepts either a plain AOBaseConfig or an FqnToConfig mapping that targets specific fully-qualified module names, checked from line 336 onward. For a plain config it looks up the handler by type in the registry dict at line 359, then calls _replace_with_custom_fn_if_matches_filter at line 122, which recurses named_modules and replaces each matching linear in place, optionally moving modules to a device first so quantization runs faster. The registry itself, declared at line 13 of the 52-line transform_module.py, is populated by register_quantize_module_handler at line 19. Prototype packages reuse the same mechanism: SmoothQuant’s api registers against the handler dict at line 151 of its module, and AWQ does the same from its core file, which is why experimental flows feel identical to stable ones.

The tensor layer is where the numerics live. quant_primitives.py, at 2388 lines, implements the reference math: quantize_affine at line 364 plus choose_qparams and dequantize primitives that do scaled integer arithmetic on plain tensors. Each subclass then specializes dispatch. Int4Tensor at line 57 comes in packing flavors, tile_packed_to_4d in a 391-line file that reshapes packed weights into four dimensions so torch.compile fuses the matmul, and a preshuffled layout matching cuBLASLt operand order. NF4Tensor at line 580 of a 1221-line file is the QLoRA substrate: to_nf4 at line 1077 builds it with 64-wide quantization blocks and 256-wide scaler blocks, and the normal-float lookup table keeps gradient flow intact for adapter training. Float8Tensor at line 84 handles the inference side, including a scaled conv3d dispatch at line 528.

The float8 training stack is a separate, older flow that predates the config registry. Its config module is 381 lines, defining the ScalingType enum at 19, Float8LinearRecipeName at 141, and Float8LinearConfig at 180, which records per-operand cast and scaling choices. convert_to_float8_training at line 86 of a 215-line utils file walks the model and swaps every nn.Linear for Float8Linear, defined at line 206, which tracks amax history and applies delayed scaling; the Float8TrainingTensor at line 240 of 376 lines carries data plus amax so FSDP2 can shard it, and float8_ops.py provides 605 lines of casting and scaled_mm helpers. This is the flow TorchTitan uses for its float8 pre-training recipes. Alongside it, the MoE prototype pairs a 1088-line mxfp8_grouped_mm.py with a 252-line config, expert-parallel dispatch and combine modules under its ep folder, and CuteDSL kernels for fused quantization.

The remaining corners round out the lifecycle. The optim package keeps optimizer states small: the 430-line adam.py implements AdamW8bit, AdamW4bit, and AdamWFp8 on one base class at line 19, shrinking state memory by 2 to 4x with only a few hundred lines of code, while CPUOffloadOptimizer at line 18 of a 180-line file pins parameters and states in CPU memory and runs fused steps there, cutting GPU VRAM by around 60 percent in the README’s example. Sparsity has a native path: ops.py, at 409 lines, registers torch.ops.torchao kernels including to_sparse_semi_structured_cutlass_sm9x_f8 at line 198, backed by CUDA sources under torchao/csrc for 2:4 semi-structured masks that gave 1.1x inference and 1.3x training speedups on SAM and ViT. And QAT’s internals are a three-file story: FakeQuantizeConfig at line 413 of fake_quantize_config.py, a 477-line fake_quantizer.py with learnable-scale variants, and a 653-line qat linear swap-in, coordinated by a 28-line two-step module that wires prepare before training and convert after.

From Install to Quantized Inference

Install the stable wheel, optionally adding the MSLK kernel package for float8 and nvfp4 inference kernels:

pip install torchao
# optional accelerated kernels
pip install mslk --index-url https://download.pytorch.org/whl/cu130

Quantize a loaded model to int4 with the tile-packed layout and the HQQ qparams algorithm, then load a float8 rowwise model through HuggingFace Transformers:

import torch
from transformers import TorchAoConfig, AutoModelForCausalLM
from torchao.quantization import (
    Int4WeightOnlyConfig,
    Float8DynamicActivationFloat8WeightConfig,
    PerRow,
    quantize_,
)

# int4 weight-only, packed for torch.compile
quantize_(model, Int4WeightOnlyConfig(group_size=32,
    int4_packing_format="tile_packed_to_4d",
    int4_choose_qparams_algorithm="hqq"))

# or load pre-quantized float8 rowwise from the Hub
config = TorchAoConfig(quant_type=Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()))
fp8_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B", dtype="auto", device_map="auto",
    quantization_config=config)

For serving, the Hub-hosted pytorch/Qwen3-32B-FP8 checkpoint runs directly under vLLM with the torchao backend enabled. For training-side use, convert_to_float8_training swaps linears before you wrap the model in FSDP2, and the QAT flow wraps any base config: pass QATConfig with a step flag of prepare to insert fake quantizers, train normally, then convert to produce the final quantized model.

Honest limits: the int4 and int8 inference kernels are tuned for linear layers on recent NVIDIA hardware, so exotic architectures, CPU-only serving, or non-linear-heavy models see smaller gains, and the ARM low-bit kernels target a narrower op set. The float8 training flow assumes accelerator support for float8 matmuls and interacts delicately with gradient accumulation and distributed setups, which is why the docs carry a compatibility table for torch and CUDA versions. MoE MXFP8 training is explicitly a prototype under torchao/prototype, and sparsity paths depend on specific tensor-core generations. None of that dims the core achievement: a quantization stack where every workflow is a config object, every quantized tensor is an honest torch.Tensor, and the whole thing compiles.

Enjoyed this post? Never miss out on future posts by following us