GuidesBackends

TorchAO Backend Guide

View as Markdown

The TorchAO backend leverages PyTorch’s torchao library for quantization-based model tuning. It provides various quantization schemes for weight-only and dynamic quantization.

Overview

  • Weight-Only Quantization: INT8, FP8
  • Dynamic Quantization: INT8, FP8, MXFP8, and NVFP4 with dynamic activations
  • Easy Configuration: Predefined quantization types
  • Pure PyTorch: No external dependencies beyond torchao

Quick Start

1from aitune.torch.backend import TorchAOBackend, TorchAOBackendConfig
2import aitune.torch as ait
3
4# Configure with FP8 weight-only quantization
5config = TorchAOBackendConfig(quantization="fp8wo")
6backend = TorchAOBackend(config)
7
8# Use in tuning
9strategy = ait.OneBackendStrategy(backend=backend)
10
11model = ait.Module(model, "my-model", strategy=strategy)
12ait.tune(model, input_data)

Quantization Types

Weight-Only Quantization

1# INT8 weight-only
2config = TorchAOBackendConfig(quantization="int8wo")
3
4# FP8 weight-only (default)
5config = TorchAOBackendConfig(quantization="fp8wo")

Dynamic Quantization

1# INT8 dynamic (activations + weights)
2config = TorchAOBackendConfig(quantization="int8dq")
3
4# FP8 dynamic (activations + weights)
5config = TorchAOBackendConfig(quantization="fp8dq")
6
7# Blackwell/Hopper-dependent dynamic quantization presets
8config = TorchAOBackendConfig(quantization="mxfp8dq")
9config = TorchAOBackendConfig(quantization="nvfp4dq")

Configuration Options

TorchAOBackendConfig

1@dataclass
2class TorchAOBackendConfig(BackendConfig):
3    fullgraph: bool = False
4    dynamic: bool | None = None
5    mode: TorchCompileMode | None = "max-autotune"
6    quantization: Literal["int8wo", "int8dq", "fp8wo", "fp8dq", "mxfp8dq", "nvfp4dq"] | None = None
7    quantization_config: AOBaseConfig | None = None
8    filter_fn: Callable[[nn.Module, str], bool] | None = None

Using Predefined Types

1config = TorchAOBackendConfig(
2    quantization="int8wo",  # Choose quantization type
3)

Custom Configuration

1from torchao.quantization import Int8WeightOnlyConfig
2
3custom_config = Int8WeightOnlyConfig()
4
5config = TorchAOBackendConfig(
6    quantization_config=custom_config,
7)

Use either quantization or quantization_config, not both.

torch.compile Options

TorchAOBackend quantizes the module and then runs it through torch.compile.

1config = TorchAOBackendConfig(
2    quantization="fp8wo",
3    fullgraph=True,
4    dynamic=None,  # None lets AITune resolve this from graph metadata
5    mode="max-autotune",
6)

Supported mode values follow torch.compile: "default", "reduce-overhead", "max-autotune", and "max-autotune-no-cudagraphs".

Filtering Modules

Use filter_fn to restrict quantization to compatible submodules. The predicate receives (module, fqn) and should return True for modules that TorchAO should quantize.

1import torch
2
3
4def linear_only(module: torch.nn.Module, fqn: str) -> bool:
5    return isinstance(module, torch.nn.Linear) and "embed" not in fqn
6
7
8config = TorchAOBackendConfig(
9    quantization="fp8dq",
10    filter_fn=linear_only,
11)

Quantization Comparison

TypeWeightsActivationsMemory ReductionSpeedAccuracy
int8woINT8FP16/FP32~2xHighBetter
int8dqINT8INT8~2xVery HighGood
fp8woFP8FP16/FP32~2xVery HighExcellent
fp8dqFP8FP8~2xVery HighExcellent
mxfp8dqMXFP8MXFP8~2xVery HighExcellent
nvfp4dqNVFP4NVFP4~4xVery HighModel-dependent

mxfp8dq and nvfp4dq require hardware and torchao support for the corresponding formats. AITune defers that validation until backend build time.

Best Practices

  1. Start with FP8: Best accuracy/performance trade-off
  2. Use INT8 for Memory: When memory is critical
  3. Dynamic Quantization: Better accuracy, slightly higher overhead
  4. Validate Accuracy: Always test quantized model accuracy
  5. Calibration Data: Use representative samples

Troubleshooting

Issue: Accuracy loss too high

Solution: Try less aggressive quantization:

1# Instead of int8wo, try fp8wo
2config = TorchAOBackendConfig(quantization="fp8wo")

Issue: Not enough speed improvement

Solution: Try dynamic quantization:

1config = TorchAOBackendConfig(quantization="fp8dq")

Next Steps