Quantization

Quantization in deep learning is the process of reducing the precision of numbers used to represent a model’s weights and/or activations. Instead of using higher-precision floating-point formats like 32-bit floats (float32) or 16-bit brain floats (bfloat16), quantization maps these values to lower-precision numerical formats, most commonly 8-bit integers (int8) or floats (fp8).

MaxText supports quantization via the Qwix library. Accurate Quantized Training (AQT) is deprecated and will be removed in a future release. Qwix is the recommended approach, providing a non-intrusive way to apply Quantized Training (QT).

Why use quantization?

The drive to use lower-precision formats like int8 or fp8 stems from significant performance advantages:

Faster computation: Hardware accelerators like TPUs and GPUs often have specialized instructions for integer arithmetic. Operations on lower-precision data like int8 or fp8 can be significantly faster than on BF16 or FP32. For example, matrix multiplications with these formats can often be 2x or more faster on hardware supporting native lower-precision tensor cores.

Reduced memory footprint: Storing weights and activations in int8 or fp8 requires 2x less memory compared to bfloat16. This reduces:

• HBM usage: Less memory is needed on the accelerator itself.

• Communication costs: Less data needs to be transferred between memory and compute units, or across devices in distributed training, which makes these transfers faster and consumes less bandwidth.

• Reduced power consumption: Lower precision operations and reduced memory access lead to less energy usage, which is crucial for deploying models on edge devices and for sustainable AI.

The primary trade-off with quantization is between the model accuracy and computational performance:

• Reduced Dynamic Range & Precision: Lower-precision formats like int8 or fp8 can represent a much smaller range of values and with less precision than BF16. This can be problematic for models with wide distributions of weights or activations, potentially clipping large values or losing fine-grained details.

• Impact on Gradients: Gradients during backpropagation can have very different, often wider, distributions than weights or activations, making them more sensitive to quantization errors.

• Convergence Issues: The approximations introduced by quantization can sometimes hinder the model’s ability to converge during training.

To overcome the challenges of quantization, libraries like Google’s Accurate Quantized Training (AQT, deprecated) and its successor Qwix (used in MaxText) employ a suite of advanced techniques. These methods ensure that models can be trained with low-precision arithmetic without significant loss in accuracy and with stable convergence.

How Quantized Training (QT) works with Qwix

Quantized Training (QT) incorporates the effects of quantization into the training loop. This allows the model to learn and adapt to the reduced precision of quantized weights and activations.

Here’s how it works:

1. Forward Pass: During the forward pass, high-precision weights and activations are converted to a lower-precision format. This step simulates the information loss that occurs during quantization. The model then performs its computations using these lower-precision representations before they are converted back to a higher precision for the rest of the network. This process forces the model to become robust to the noise and reduced range of quantized values.

2. Backward Pass: Standard backpropagation cannot flow through the non-differentiable quantization operations (like rounding). To solve this, QT uses the Straight-Through Estimator (STE). The STE essentially “ignores” the non-differentiable quantization step during the backward pass, passing the gradients through as if the operation was an identity function. This allows the high-precision weights to be updated based on the loss, enabling the model to learn effectively.

By integrating the quantization simulation directly into the training, the model learns to minimize the impact of precision loss, resulting in a more accurate quantized model.

Using Quantization in MaxText

You can enable quantization in MaxText by setting flags in your configuration file (e.g., base.yml) or via the command line. MaxText supports Qwix (recommended) and the legacy AQT library (deprecated).

Configuration Flags

The primary flags to control quantization are:

• use_qwix_quantization: A boolean flag.

  • Set to True to enable quantization using the Qwix library.

  • Set to False (or omit) to use the AQT library (deprecated) if quantization is set.

• quantization: A string that specifies the type of quantization to apply. The accepted values depend on whether you are using Qwix or legacy AQT.

• quantization_calibration_method: The calibration method for weights and activations (e.g., "absmax"). This is mainly for Qwix.

• quantize_router_proj (default: True for backward compatibility): Applicable when use_qwix_quantization=True and quantization is set; ignored otherwise. (1) If enabled, quantizes the MoE router (gate) projection matmul. (2) Set to False to keep the router projection unquantized across all quantization schemes (e.g., fp8_full, int8).

  • Enabling this together with float32_gate_logits=True is rejected at config init, because the fp32 cast on the gate operands would be undone by requantization at the projection matmul — see MoE configuration for the full rationale.

• quantize_mtp (default: False): Applicable when use_qwix_quantization=True and quantization is set; ignored otherwise. If enabled, quantizes the Multi-Token Prediction (MTP) block. Requires mtp_num_layers > 0 and quantization='fp8_full'.

• quantize_logits_proj (default: False): Applicable when use_qwix_quantization=True and quantization is set; ignored otherwise. If enabled, quantizes the output logits (logits_dense) projection matmul. Requires logits_via_embedding=False (un-tied embeddings) and quantization='fp8_full'.

  • logits_proj_quant_calibration_method (default: ""): Calibration method for the logits projection. If empty (default), inherits weight_quantization_calibration_method and act_quantization_calibration_method. Set to e.g. 'absmax' to force absmax calibration.

  • logits_proj_bwd_quant_calibration_method (default: ""): Backward (gradient) calibration method for the logits projection, applied to both the activation-gradient and weight-gradient arms. If empty (default), inherits bwd_quantization_calibration_method and drhs_grad_quantization_calibration_method. Set to e.g. 'absmax' to force absmax calibration.

  • Enabling this together with logits_dot_in_fp32=True is rejected at config init, because the fp32 cast on the logits operands would be undone by requantization at the projection matmul.

  • MTP has no output head of its own — it reuses the shared logits_dense via apply_output_head. This flag therefore also quantizes the MTP logits projection, independently of quantize_mtp.

Qwix Quantization (Recommended)

To use Qwix, you must set use_qwix_quantization=True. Qwix is a powerful and non-intrusive library for Quantized Training.

quantization values for Qwix

Common options for the quantization flag when using Qwix include:

• "int8": 8-bit integer quantization.

• "fp8": 8-bit floating-point quantization.

• "fp8_full": FP8 quantization with static scaling.

• "fp8_gpu": FP8 for NVIDIA GPUs.

• "fp8_nanoo": FP8 for AMD MI300/MI325 GPUs.

Example command for Qwix

Here is an example of how to run a training job with int8 quantization enabled via Qwix:

python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?} base_output_directory=gs://<GCS_BUCKET> dataset_type=synthetic use_qwix_quantization=true quantization='int8'

The Qwix Interception API

MaxText integrates Qwix using its powerful and non-intrusive Interception API. This approach allows you to enable QAT for your models without modifying the original model source code. You don’t need to manually replace nn.Dense with QuantizedDense or other quantized layer types.

Instead, you define a set of quantization rules externally. Qwix then uses a context manager to “intercept” the creation of standard Flax/NNX layers during model initialization and dynamically replaces the layers with their QAT-enabled versions on the fly.

A quantization rule can be defined as follows:

rule = [
    qwix.QtRule(
        module_path="decoder/.*layers.*",
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        bwd_qtype=jnp.int8,
        bwd_weight_grad_tile_size=1 / config.quantization_local_shard_count,
        op_names=("dot_general",),
    )
]

QtRule parameters:

• module_path: A regex to match the layers to which this rule should be applied.

• weight_qtype: The target quantization type for weights (e.g., jnp.int8).

• act_qtype: The target quantization type for activations.

• bwd_qtype: The quantization type for the backward pass.

• op_names: The operations to be quantized (e.g., "dot_general").

This rule is then used within a QtProvider to quantize the model automatically:

model = qwix.quantize_model(model, qwix.QtProvider(rule))

AQT Quantization

[!WARNING] DEPRECATION NOTICE: AQT quantization is deprecated and will be removed in a future release. Please migrate to Qwix by setting use_qwix_quantization=True.

If use_qwix_quantization is False or not set, you can still apply quantization using the AQT library by setting the quantization flag. You can read more about AQT on this Google Cloud blog.

quantization values for AQT

When using AQT, you can pass one of the following values to the quantization flag:

• ‘int8’ for dynamic range quantization using 8-bits

• ‘int8w’ for weights only quantization using 8-bits

• ‘int4w’ for weights only quantization using 4-bits

• ‘intmp’ for mixed precision weight only quantization based on config file

• ‘fp8’ for 8-bit floating-point GeMMs on NVIDIA GPUs.

Example command for AQT

python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?} base_output_directory=gs://<GCS_BUCKET> dataset_type=synthetic use_qwix_quantization=false quantization='int8'

Note that use_qwix_quantization is not set to True.

For further reading, please refer to the Qwix Read the Docs website.

DeepSeek V3 Fine-tuning FP8 Recipe

To improve the performance of DeepSeek V3 fine-tuning, we developed a custom recipe optimized for FP8 throughput. The method prioritizes specific compute-intensive and bandwidth-heavy components while preserving training stability through a fine-grained scaling strategy.

Quantization Scope

To realize these gains, the recipe employs a w8a8g8 (8-bit weights, activations and gradients) strategy targeting three primary areas:

• Megablox Kernels: Specifically the gmm and tgmm operations.

• Attention Projections: Utilizing convolution fusion.

• Communication: Specifically the weight All-Gathers.

FP8 Recipe

• Rounding: rounding to nearest even

• Precision

  • Activations and weights: e4m3fn

  • Gradients: e5m2

• Scaling granularity: per-axis

• Scaling mode:

  • static for weights and activations

  • dynamic for gradients

Convergence

To validate this recipe, we utilized MaxText following the MLPerf Training framework by MLCommons to ensure a reproducible and standardized evaluation. Using the C4 dataset (loaded via TFDS) as the reference corpus, we tracked convergence by monitoring validation loss on a held-out split. This aligns with MLPerf’s time-to-quality principle, where the primary metric is the speed at which the model achieves target quality.

For this specific case, we derived our training duration from the MLPerf 405B benchmark, targeting roughly 2–3 billion tokens after resuming from a checkpoint. In our configuration, we executed 300 steps with a sequence length of 4096 and a global batch size of 2048, resulting in a total of approximately 2.5 billion tokens.

Performance Sensitivity

Please note that the FP8 benefits are highly sensitive to model parameters, the efficiency of the BF16 baseline, and hardware utilization; consequently, results will vary when this recipe is applied to other models. Any variance in these factors shifts the ratio of compute-bound to memory-bound operations, directly altering the potential gains.

Quantized Weight All-Gather (Weight QAG) for GMM

In MoE models, expert weights represent a substantial memory and communication footprint under Fully Sharded Data Parallelism (FSDP). When expert weights are sharded across FSDP ranks, MaxText supports Quantized Weight All-Gather (Weight QAG) to quantize weights before gathering them across the FSDP mesh axis and passing to the Grouped Matrix Multiplication (GMM) kernel, reducing inter-chip weight communication volume proportionally to the quantized width (2x for FP8 weights gathered in place of BF16).

quantize weights -> FSDP all-gather -> GMM

Mechanism

1. Pre-Gather Quantization: Each FSDP shard quantizes its local weight slice from the model dtype (e.g. BF16) to the quantized dtype (e.g. FP8).

2. FSDP All-Gather: Gathers the quantized byte payloads across the FSDP axis, shrinking communication volume by the compression ratio compared to gathering in the model dtype.

3. GMM Forward: The gathered quantized weights are consumed directly by the grouped GEMM kernel without re-quantization.

4. Backward Pass: In the backward pass, gradients w.r.t. weights (drhs) are computed and reduce-scattered back across the FSDP ranks.

Usage & Requirements

Weight QAG is triggered when an MoE FSDP sharding strategy is combined with quantization.

Requirements:

• Either shard_embed_moe_on_fsdp=True or shard_exp_on_fsdp=True. See MoE Configuration.

• quantization must be specified (e.g. fp8_full).

• Static weight calibration: weight_quantization_calibration_method starting with fixed (ensures uniform scale factors across shards without collective synchronization).

Token Quantized All-Gather Pipeline for MoE Ring of Experts

In Expert Parallelism (EP) sharding and Ring of Experts (use_ring_of_experts=True), MaxText provides the option to quantize token activations prior to the All-Gather across the EP mesh axis, reducing inter-chip communication volume proportionally to the quantized width (2x for FP8 activations gathered in place of BF16).

Along with quantized token all-gather, MaxText defines the entire quantized token execution pipeline to reduce re-quantization overhead:

quantize token -> EP all-gather -> sort -> GMM

Quantization Pipeline

1. Quantize Token: Activations are quantized from the model dtype (e.g. BF16) to the quantized dtype (e.g. FP8).

2. EP All-Gather: Gathers the quantized byte payloads across the EP mesh axis, shrinking inter-chip bandwidth by the compression ratio.

3. Sort: Token sorting (default or ragged sort) operates directly on the raw quantized payloads without dequantizing to float.

4. GMM: Tokamax GMM v2 computes directly on the quantized activations and weights, and scales accumulated outputs back to model dtype.

5. Backward TGMM: Receives cached pre-quantized activations directly from the forward pass, avoiding re-quantization.

Usage & Requirements

Flag: moe_quantize_token_all_gather=true

Requirements:

• use_ring_of_experts=True: Inherent architectural requirements.

• use_gmm_v2=True: Tokamax GMM v2 grouped GEMM backend.

• Static calibration for activations: e.g. act_quantization_calibration_method="fixed,-224,224" (dynamic calibration is for future work).