DOCSEARCHbyAlgolia
Submit Docs

AQT (Accurate Quantized Training)

repository·main·Indexed 18 days ago

https://github.com/google/aqt

A JAX-based software library for tensor operation quantization, designed to enable high-quality int8 quantization in both training and serving to minimize training-serving bias. AQT provides low-level quantization primitives for JAX (such as quantized_dot and quantized_sum) and high-level Flax modules including DenseAqt, ConvAqt, EmbedAqt, LayerNormAqt, and MultiHeadDotProductAttentionAqt.

Tokens
10.9K
Snippets
25
Records
46
Agent score
60%

What's inside AQT

  1. Overview of AQT (Accurate Quantized Training)

    main
    AQT (Accurate Quantized Training) provides quantization for convolution and matmul operations, following the principle of what you serve is what you train. It provides specialized libraries for both JAX and Flax to enable quantized training and inference.

  2. Overview of AQT Quantization Library

    main
    The AQT quantization libraries for Jax and Flax provide "what you serve is what you train" quantization for convolution and matmul operations. For detailed usage of the quantization API, refer to the main JAX documentation in the jax/ directory.

  3. Overview of ImageNet classification with AQT

    main
    This module provides a ResNet50 implementation for the ImageNet classification task. It supports optional quantization of weights and activation functions for Matmul layers, including all Conv layers and the Dense layer. The training process uses linear learning rate warmup and a cosine learning rate schedule.

  5. Understand Leaderboard Transformer Quantization configurations

    main

    The leaderboard configurations summarize quantization results for various transformer model setups. These configurations include:

    • Baseline models: Running in bfloat16.
    • Weight-only quantized models: Supports 8bit, 4bit, and 2bit quantization.
    • Weight & Activation quantized models: Supports 8bit and 4bit quantization.

    For models utilizing activation quantization, experiments are conducted using two different bounding strategies:

    1. Fixed bounds: Pre-defined ranges for activations.
    2. Dynamic bounds: Bounds that are automatically adjusted during the training process.

  6. Understand AQT INT8 internal mechanics

    main

    AQT achieves high quality by using per-example (for activations) and per-output-channel (for weights) scaling. This isolates the impact of outliers to a single row or column, allowing for tighter calibration.

    Internally, a simple INT8 AQT matmul follows these steps:

    1. Quantization: Clips and rounds values to the INT8 range (e.g., -127 to 127).
    2. Calibration: Calculates scales ($a_s$ and $w_s$) based on the maximum absolute values of the inputs.
    3. Computation: Performs an integer matmul with an int32 accumulator, then de-quantizes the result using the scales.

    To get hardware acceleration (like on TPUs) in JAX, AQT relies on jnp.matmul with preferred_element_type=jnp.int32, which maps to lax.dot_general and XLA's DotGeneral op.

    import jax.numpy as jnp

def matmul_true_int8(lhs, rhs):
  assert lhs.dtype == jnp.int8
  assert rhs.dtype == jnp.int8
  # Uses int32 accumulator for hardware acceleration
  result = jnp.matmul(lhs, rhs, preferred_element_type=jnp.int32)
  assert result.dtype == jnp.int32
  return result

def aqt_matmul_int8(a, w):
  max_int8 = 127
  def quant_int8(x):
    return jnp.clip(jnp.round(x), -max_int8, max_int8).astype(jnp.int8)

  # Calibration: per-example and per-channel scales
  a_s = max_int8 / jnp.max(jnp.abs(a), axis=1, keepdims=True)
  w_s = max_int8 / jnp.max(jnp.abs(w), axis=0, keepdims=True)

  # int8 matmul with int32 accumulator, then de-quantize
  result = matmul_true_int8(quant_int8(a * a_s), quant_int8(w * w_s)) / (a_s * w_s)
  return result

  7. How AQT handles weight transformations before matmul

    main

    In many models, a transformation $T$ is applied to weights before a matmul (e.g., matmul(a, T(w))).

    Standard quantization might compute matmul(a, T(Q(w))), which saves checkpoint size but doesn't accelerate the matmul because $T$ likely returns floating-point values.

    To achieve both checkpoint compression AND hardware acceleration, AQT aims to store the fully transformed and quantized weight $w_q = Q(T(w))$ in the checkpoint. This allows serving to use matmul(a, w_q) directly.

    AQT provides a specialized matmul_aqt function that handles this: matmul_aqt(..., T(w)) = matmul(..., Q(T(w))). Note that $Q(T(w))$ is not visible outside of matmul_aqt because the function uses a custom gradient definition.

  8. How Intercept Methods differ from Injection API

    main

    AQT provides two primary ways to apply quantization to Flax modules:

    1. Injection API: Requires placing quantization configurations directly inside the module's code.
    2. Intercept Methods API: Uses flax.linen.intercept_methods to intercept Flax module methods at runtime. This allows you to apply quantization by execution scope via a context manager, leaving the original module code untouched.

    Functionally, the Intercept API is designed to produce the same results as the Injection API, but offers better usability for existing models where you do not want to refactor the source code.

  9. Quantization support for Machine Translation

    main

    In the WMT Machine Translation example, AQT-style quantization is supported for all Matmul layers. This includes:

    • All DenseGeneral layers
    • The Dense layer

    Quantization configurations (such as 4-bit weights and auto-activation) are typically passed via the --hparams_config_dict flag in the train.py script.

  10. How quantization injection works in AQT

    main
    AQT uses a technique called "quantization injection." In JAX-based neural network libraries, tensor contraction operations (like jax.numpy.einsum or flax.linen.DenseGeneral) call lax.dot_general as their core computation. To quantize a model, you substitute lax.dot_general with a quantized variant provided by AQT. JAX-based libraries like Flax and Pax provide APIs to perform this substitution when creating layers.

  11. Install the AQT library via pip

    main

    To install the AQT package with support for JAX legacy components, use pip. It is recommended to upgrade pip before installation.

    # Upgrade pip.
pip install --user --upgrade pip

# Install AQT package
pip install aqtp[jax_legacy]

  12. Generate experiment reports with the Reporting Tool

    main
    Once a training run is complete, you can use the reporting tool located in report_utils.py to generate concise experiment reports containing aggregated metrics and metadata. Detailed instructions can be found in the utils/ directory documentation.