TorchAO Quantization¶

Overview¶

TorchAO provides quantization tools for PyTorch and supports both pre-quantized checkpoint loading and runtime quantization.

Hardware Support¶

Legend: ✅ supported, ❌ unsupported, ⭕ not verified in this guide.

Model Type Support¶

Diffusion Model (Boogu-Image)¶

Configuration¶

Using Boogu-Image as an example:

vllm serve Boogu/Boogu-Image-0.1-Base-fp8 \
  --omni \
  --port 8091 \
  --diffusion-quantization-config \
  '{"transformer":{"method":"torchao_float8_weight_only"}}'

The equivalent command using the complete serialized TorchAO configuration is:

vllm serve Boogu/Boogu-Image-0.1-Base-fp8 \
  --omni \
  --port 8091 \
  --diffusion-quantization-config \
  '{
    "transformer": {
      "method": "torchao",
      "quant_type": {
        "default": {
          "_type": "Float8WeightOnlyConfig",
          "_version": 2,
          "_data": {
            "weight_dtype": {
              "_type": "torch.dtype",
              "_data": "float8_e4m3fn"
            },
            "set_inductor_config": false
          }
        }
      }
    }
  }'

Parameters¶

See TorchAO's vLLM integration guide for the serialized form and workflow configs for available configuration classes and parameters.

Validation and Notes¶

At load time, vLLM-Omni builds a TorchAOConfig from --diffusion-quantization-config. For the Boogu-Image checkpoints documented here, indexed PyTorch .bin shards are loaded through pt_weights_iterator.

When a component map is used, only the components included in the configuration use TorchAO. Other components keep their own checkpoint and runtime settings.

Use either the torchao_float8_weight_only shorthand or the equivalent full quant_type configuration shown above. The checkpoint must already contain weights quantized with TorchAO.