---
title: 'TorchAO & Brevitas: Quantization & Sparsity'
url: https://www.emergentmind.com/topics/torchao-and-brevitas
type: topic
---

# TorchAO & Brevitas: Quantization & Sparsity

TorchAO and Brevitas are quantization and sparsity frameworks for deep learning model optimization, focused on compressing, accelerating, and enabling efficient deployment of large language models (LLMs) and similar architectures. TorchAO is a PyTorch-native, end-to-end solution covering pre-training through serving, tightly integrated with PyTorch APIs and ecosystem tools, offering innovations in low-precision tensor representations, sparsity, and workflow orchestration. Brevitas, in contrast, was an early quantization-aware training solution for PyTorch, supporting INT8/INT4 workflows but lacking many of TorchAO’s extensibility and workflow features. This article presents a rigorous account of TorchAO’s architecture, quantization and sparsity techniques, performance characteristics, and an in-depth comparison to Brevitas, based exclusively on reported details in arXiv:2507.16099.

## 1. Architectural Overview and Workflow Integration

TorchAO implements a native, unified training-to-serving pipeline for LLM optimization, integrating quantization and sparsity at every model lifecycle stage [2507.16099]. Its architecture organizes optimization as modular steps for pre-training, fine-tuning, and inference:

- **Pre-Training:** FP8 low-precision training is enabled via calls to `torch.compile(model)` and `convert_to_float8_training(model)`. TorchTitan manages distributed training across FSDP2, tensor parallelism, and asynchronous tensor parallel backends, leveraging FP8 GEMM kernels on supported hardware.
- **Fine-Tuning:** Quantization-Aware Training (QAT) is instrumented by calling `quantize_(model, IntXQuantizationAwareTrainingConfig(...))` to insert fake quantization operations. After training, quantized ops are materialized for deployment using `quantize_(model, FromIntXQuantizationAwareTrainingConfig())`. Tools such as TorchTune and Axolotl are integrated at this stage.
- **Serving:** Post-Training Quantization (PTQ) is performed by `quantize_(model, <PTQConfig>)`. TorchAO supports server-side backends (HuggingFace Transformers via `TorchAoConfig`, vLLM, SGLang with GemLite Triton kernels) and edge deployment through ExecuTorch exports to XNNPACK and custom sub-8-bit ARM/Metal kernels.

The model dataflow is defined as:

```
Model definition → (optional FP8/QAT) → checkpoint → (optional PTQ/sparsity) → export → backend runtime
```

One-line transformations, as exemplified by
```python
from torchao.quantization import Int4WeightOnlyConfig
from transformers import AutoModelForCausalLM, TorchAoConfig
cfg = TorchAoConfig(quant_type=Int4WeightOnlyConfig())
model = AutoModelForCausalLM.from_pretrained(..., quantization_config=cfg)
```
facilitate in-place quantization and backend transition. TorchAO thus unifies an otherwise fragmented space across model development, optimization, and deployment [2507.16099].

## 2. Quantization Techniques and Supported Data Types

TorchAO supports both QAT and PTQ, including training and inference in novel low-precision formats beyond common INT8 and INT4 [2507.16099].

**Quantization Mapping:** For tensor element $x$, quantization to lower precision is given by
$$
q(x) = \mathrm{clip}\bigl(\round(x/s), q_{\min}, q_{\max}\bigr)\times s
$$
where $s$ is a learned/computed scale; $q_{\min},q_{\max}$ are format-specific (e.g., INT4: $[-8,7]$, INT8: $[-128,127]$). FP8 utilizes a custom 4-bit mantissa and 3-bit exponent plus sign.

**Supported Data Types:**
- Integer: INT4, INT8
- Floating: FP8 (rowwise, tensorwise), MXFP4, MXFP6, MXFP8 (prototypes)
- Mixed-precision: NF4 (for QLoRA)

**PTQ and QAT APIs:** PTQ is performed by calls such as `quantize_(model, Int8DynamicActivationInt4Weight(group_size=32))`. QAT workflows leverage:
```python
activation_cfg = FakeQuantizeConfig(torch.int8, "per_tensor", False)
weight_cfg = FakeQuantizeConfig(torch.int4, group_size=32)
quantize_(model, IntXQuantizationAwareTrainingConfig(activation_cfg, weight_cfg))
# After training:
quantize_(model, FromIntXQuantizationAwareTrainingConfig())
```
Group size, per-tensor and per-channel scaling, and flexible hardware-aligned quantization schemas are supported.

TorchAO supports QAT and PTQ for INT4/8, FP8, and experimental MXFP types, with tunable group sizes, dynamic and static quantization, and hardware-targeted quantization for efficient inference [2507.16099].

## 3. Sparsity: Techniques and Implementation

TorchAO provides inbuilt sparsity support, primarily implementing 2:4 sparsity, where each four-element group contains exactly two zeros:
Given weight vector $w$, partition $w$ into blocks $G_k$ of size 4; for each block require
$$
\#\{\,i\in G_k : w_i =0\} = 2,\;\forall k
$$
APIs include `sparsify_(model, SemiSparseWeightConfig())` for 2:4, or `sparsify_(model, BlockSparseWeightConfig())` for block patterns. Internally, sparse layouts are encoded via tensor subclasses that store block masks and quantized values together, allowing fusion with backend sparse GEMM kernels (notably for Tensor Core 2:4 support). *A plausible implication is* that these internal representations facilitate efficient, hardware-aligned execution without requiring model source code changes [2507.16099].

## 4. Tensor Subclass Abstraction: QuantizedTensor and Modularity

A foundational innovation of TorchAO is the use of torch.Tensor subclassing, allowing creation of QuantizedTensor objects that encode all quantization metadata—scale, zero_point, bit_width, and layout—while supporting autograd and distributed semantics. In the forward pass, modules such as FakeQuantizedLinear wrap standard nn.Linear layers and produce QuantizedTensors as output for activations and weights.

Upon conversion for backend execution, QuantizedTensor buffers can be directly passed to CUDA, ARM, or Metal kernels, with operation dispatch managed by the `__torch_dispatch__` override. This model permits hardware-agnostic quantization and sparsity, seamless interoperability with PyTorch model APIs, and avoids intrusive graph rewrites. A simplified class hierarchy is:

| Parent Class  | Subclass          | Key Attributes                  |
|---------------|-------------------|---------------------------------|
| `torch.Tensor`| QuantizedTensor   | scale, zero_point, bit_width    |

This abstraction supports lossless transitions between fake and real quantization during QAT and PTQ, and underpins the extensibility of TorchAO’s quantization kernel interface [2507.16099].

## 5. Performance Characterization and Evaluation

TorchAO provides empirical results demonstrating low-precision model optimization benefits for Llama3-8B and related models [2507.16099].

**FP8 Training (Table 1):**
| Scaling         | Peak Mem (GB) | Median tok/s | Speedup  |
|-----------------|---------------|--------------|----------|
| BF16 (baseline) | 47.65         | 6150         | 1.00×    |
| tensorwise+FP8  | 47.77         | 7689         | 1.25×    |

**QAT on Llama3-8B (Table 2):**
| Model           | quant hellaswag | quant wikitext ppl | train tok/s | peak mem (GB) |
|-----------------|-----------------|--------------------|-------------|--------------|
| PTQ             | 47.0%           | 26.27              | 480.3       | 17.6         |
| QAT             | 52.8% (+69.8%)  | 12.31 (+82.8%)     | 323.0 (−33%)| 32.9 (+87%)  |

**PTQ for Inference (Table 3):**
| Quant Technique | Accuracy | ppl   | tok/s  | size (GB) |
|-----------------|----------|-------|--------|-----------|
| None (BF16)     | 60.01    | 7.33  | 132.41 | 15.01     |
| int4wo-64       | 58.10    | 8.25  | 268.88 | 4.76      |
| float8wo        | 59.83    | 7.37  | 213.88 | 8.03      |

**Serving on vLLM (Table 4):**
| Precision      | tok/s  | per-token latency (ms) | inter-token latency (ms) |
|----------------|--------|-----------------------|--------------------------|
| BF16           | 103.6  | 9.50                  | 9.47                     |
| FP8 tensorwise | 132.8  | 7.48 (−21%)           | 7.47 (−21%)              |

These results demonstrate that FP8 and INT4/8 quantization, as well as sparsity optimizations, yield substantial improvements in inference throughput and memory efficiency, often with minimal accuracy degradation [2507.16099]. *This suggests* TorchAO’s techniques are viable for production LLM deployment.

## 6. Comparison of TorchAO and Brevitas

The following summarizes key differentiators between TorchAO and Brevitas as reported in [2507.16099]:

| Aspect                 | TorchAO                | Brevitas                                     |
|------------------------|-----------------------|----------------------------------------------|
| API Design             | 1-line in-place transforms (`quantize_`, `sparsify_`, `convert_to_float8_training`) | Explicit quantizer modules, graph rewrites   |
| HF Transformers Integration | TorchAoConfig enables seamless `save_pretrained`/`push_to_hub` | Manual export of quantized parameters        |
| Bit-Precision Flexibility   | Supports FP8, custom MXFP, NF4, INT4/8; flexible groupings/scaling | Primarily INT8/INT4 (QAT/PTQ), no FP8/hardware-tied formats |
| Backend Support        | CUDA, Triton (GemLite), XNNPACK, ARM, Metal  | CUDA, CPU reference kernels                  |
| Workflow Coverage      | Pre-training, fine-tuning, serving spanning entire LLM lifecycle | QAT/PTQ for model compression; training/serving left to user  |
| Tensor Representation  | torch.Tensor subclassing (preserves autograd, distributed semantics) | Module hooks, FakeQuantize ops               |
| Sparsity Patterns      | 2:4/block/Marlin layouts integrated in kernels | No built-in support                          |
| Prototyped Advances    | AutoRound, AWQ, ParQ, SmoothQuant, SpinQuant | Not present                                  |

TorchAO’s close integration with the PyTorch and HuggingFace ecosystems, broad data type support—including FP8 and custom formats—hardware-targeted quantization, sparsity, and extensibility in both research and production distinguish it from Brevitas. *A plausible implication is* that TorchAO serves as a preferred optimization substrate for LLM training and deployment where full-lifecycle workflow and ecosystem integration are operationally critical [2507.16099].

Source: https://www.emergentmind.com/topics/torchao-and-brevitas