Skip to main content
W&B Fully Connected
>
Articles

Quantization-Aware Training (QAT): A step-by-step guide with PyTorch

A practical deep dive into quantization-aware training, covering how it works, why it matters, and how to implement it end-to-end.
Brett Young
Created on March 25|Last edited on March 28
Quantization is the process of converting a model’s weights and activations from high-precision formats like FP32 to lower-precision formats such as INT8. This reduces the model's memory footprint, decreases bandwidth requirements, and can speed up inference—when the backend or hardware supports low-precision operations. These benefits make quantization a practical choice for deploying models on edge devices like smartphones, IoT hardware, and embedded systems.
Many modern mobile CPUs, NPUs, and specialized accelerators now include optimized support for INT8 operations, enabling faster and more efficient inference compared to full-precision models.
Among the commonly used approaches to quantization, post-training quantization (PTQ) is the simplest. It applies quantization after the model has been trained, without modifying the training pipeline. While PTQ is easy to implement and often works well for large, robust models, it can lead to noticeable accuracy drops in smaller or more sensitive architectures.
Quantization-aware training (QAT) takes a more integrated approach by simulating quantization during training itself. Fake quantization modules are added to the model so it learns to operate within the constraints of lower precision. This allows the model to adapt and usually results in better accuracy after quantization, making QAT a better option when preserving model performance is important.


Table of contents

Table of contentsWhat is quantization?QAT vs. PTQ: Which training quantization method is best?Handling non-differentiable operations in QATAdvantages and challenges of QATTutorial: Fine-tuning a quantized model using TorchaoRunning inference with the quantized model Looking ahead: What’s next in LLM quantization?Conclusion


What is quantization?

Quantization is the process of converting high-precision floating-point values—typically FP32—into lower-bit integer representations such as INT8, INT4, or even binary formats. This applies to model weights, activations, or both, depending on the setup. The goal is to reduce memory usage and improve computational efficiency, especially during inference.
In practice, quantization precisely maps continuous floating-point values (like FP32) onto a discrete integer range (such as INT8, which includes numbers from –128 to 127). To achieve this accurately, quantization relies on two parameters: the scale and the zero-point.
The scale controls how finely or coarsely the integer numbers represent the original floating-point values. A smaller scale provides finer granularity, giving better precision but limiting the range of numbers that can be represented. A larger scale covers a broader range but results in reduced precision.
The zero-point shifts the integer range so the real number zero exactly corresponds to an integer. Without this offset, representing zero would involve approximation, introducing unnecessary errors in operations like zero-padding or ReLU activations.
Once values have been quantized, neural networks typically perform calculations directly using integer arithmetic. However, at some stages—such as interpreting model outputs or transitioning between layers that require floating-point inputs—the quantized integers need to be converted back into approximate floating-point numbers. This reverse operation is called dequantization. During dequantization, the integer value is scaled and shifted back to a floating-point approximation of the original number, using the same scale and zero-point parameters defined during the quantization step.
Thus, quantization and dequantization together allow neural networks to benefit from efficient integer arithmetic while closely preserving the accuracy of computations originally performed with floating-point precision. However, not all backends support true integer execution. On some hardware where quantization kernels have not been implemented, quantized weights may be converted back to FP16 or FP32 before computation. This means the model is stored in low-bit format for memory savings, but the actual operations are performed in higher precision, often negating the speed benefits.

QAT vs. PTQ: Which training quantization method is best?

There are two common approaches to quantizing neural networks: Quantization-Aware Training (QAT) and Post-Training Quantization (PTQ). They differ mainly in when quantization is applied and how much control you have over the accuracy tradeoffs.
PTQ is the simpler option. It’s applied after the model has already been trained using full-precision (e.g., FP32) weights and activations. Once training is complete, these values are converted to lower precision formats like INT8. This process is fast, lightweight, and requires no retraining—but because the model wasn’t trained with quantization in mind, the change in precision can cause a drop in accuracy, especially for smaller or more sensitive models.
QAT, by contrast, introduces quantization into the training process itself. During QAT, the model simulates low-precision behavior using “fake quantization” modules inserted into the forward pass. These modules mimic the effects of quantization by rounding values to a discrete integer grid and then immediately converting them back to floating point using a scale and zero-point. This lets the model experience the kinds of errors it will encounter during inference—like rounding and clipping—while still allowing training to proceed using high-precision math (e.g., FP32) for gradient updates.
Because rounding is not differentiable, QAT uses a trick called the straight-through estimator (STE) to approximate gradients. In essence, STE treats non-differentiable steps like rounding as if they were identity operations during the backward pass, allowing gradients to flow through the fake quantization layers. While this isn’t a perfect representation of the underlying math, it works well in practice and enables the model to adapt to quantization noise during training.
Although QAT is more computationally expensive—it requires additional training time and more careful tuning—it often results in significantly better accuracy under quantized constraints. This is especially true when targeting lower bit-widths like INT4, or when working with models that are highly sensitive to small changes in precision, such as compact CNNs or transformers.
In practice, PTQ is often good enough for large, robust models or use cases where slight accuracy drops are acceptable. But for edge deployments, performance-critical applications, or aggressive quantization levels, QAT is typically the better choice.

Handling non-differentiable operations in QAT

A significant challenge in quantization-aware training (QAT) arises from the inherent non-differentiability of quantization operations—especially the rounding step. Neural network training depends heavily on gradient-based optimization methods, such as backpropagation, which require all operations within the computational graph to be differentiable. However, quantization involves rounding continuous values to discrete integers, an operation that inherently lacks a meaningful gradient.
To address this challenge, QAT commonly employs techniques like the straight-through estimator (STE). STE provides a practical way to approximate gradients for non-differentiable operations by treating the rounding operation as an identity function during the backward pass. In other words, although rounding occurs during the forward pass, the backward pass "ignores" the rounding and simply passes the gradient unchanged. This approximation allows gradients to flow through quantization operations, enabling gradient-based optimization to continue effectively despite the discontinuities introduced by quantization.
While STE introduces some approximation errors—because it simplifies the true behavior of the rounding operation—it is typically sufficient in practice to train quantized models effectively. More advanced techniques have been explored, including smoother approximations and learned quantization boundaries, but STE remains the most widely used and simplest approach.
Thus, by employing methods like STE, QAT circumvents the fundamental challenge posed by rounding and other non-differentiable quantization operations, enabling effective training of highly accurate quantized models.

Advantages and challenges of QAT

Quantization-Aware Training significantly benefits model performance and deployment efficiency by explicitly training neural networks to handle lower-precision arithmetic. Because QAT introduces quantization effects during training, the model learns to adjust to reduced precision, minimizing accuracy losses commonly associated with quantization.
A major benefit is that QAT allows aggressively quantized models—such as INT8 or even INT4—to retain near-original accuracy, which significantly reduces memory usage and the amount of data transferred during inference. Smaller model sizes mean faster loading and less memory required on hardware, enabling efficient deployment even on resource-constrained edge or mobile devices. Additionally, because QAT-trained models perform inference directly in low-precision arithmetic, they run faster on hardware specifically optimized for integer operations.
Moreover, QAT often leads to models that are more robust to the noise introduced by quantization, improving overall reliability in production. This is especially valuable in environments where performance and power efficiency are critical.
Although these advantages are substantial, QAT does come with challenges—such as higher computational costs during training, additional complexity in hyperparameter tuning, and the need for specialized methods to handle non-differentiable operations like rounding. However, in practice, the performance gains and efficiency improvements often outweigh these costs.

Tutorial: Fine-tuning a quantized model using Torchao

n this tutorial, we fine-tune a causal language model using quantization-aware training (QAT) with Torchao. Specifically, we use the Int8DynActInt4WeightQATQuantizer class to prepare the Qwen/Qwen2.5-0.5B model for deployment with significantly reduced precision: weights quantized to 4 bits per channel group, and activations dynamically quantized to 8 bits per token.
By simulating quantization during training, the model learns to operate effectively under these constraints—helping retain accuracy even at aggressive bit-widths.
We start by loading a pretrained model and tokenizer, then prepare the model for QAT by replacing standard linear layers with quantization-aware layers. We fine-tune the model using Hugging Face’s Trainer API and a tokenized language modeling dataset. After training, we convert the simulated layers to actual quantized modules, resulting in a compact, efficient model ready for deployment.
Here’s the full training script:

import torch
import torch.nn as nn
from torchao.quantization.prototype.qat import Int8DynActInt4WeightQATQuantizer
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
    TrainerCallback,
    set_seed
)
from datasets import load_dataset
import os
import copy
import logging
import wandb 


# Setup logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

# Set seed for reproducibility
set_seed(42)

# Model and training configuration
model_name = "Qwen/Qwen2.5-0.5B"
output_dir = "./qwen_int4_quantized"
num_train_epochs = 1
per_device_train_batch_size = 1
gradient_accumulation_steps = 16
learning_rate = 5e-5
max_seq_length = 512
warmup_steps = 100
logging_steps = 10
save_steps = 50
eval_steps = 50
save_total_limit = 2
fp16 = True  # Mixed precision training

# Quantization parameters
groupsize = 224
padding_allowed = False

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token  # Set padding token if not defined

# Make sure the model is loaded in FP32
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)

# Dataset preparation
def prepare_dataset():
    """Load and prepare the dataset for training"""
    # For this example, we'll use a small text dataset
    # Replace with your actual dataset
    
    # Option 1: Use wikitext
    dataset = load_dataset("wikitext", "wikitext-2-v1")
    
    # Use only a portion for quick testing
    train_dataset = dataset["train"].select(range(min(10000, len(dataset["train"]))))
    eval_dataset = dataset["validation"].select(range(min(100, len(dataset["validation"]))))
    
    def tokenize_function(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=max_seq_length,
            padding="max_length",
            return_tensors="pt",
        )
    
    tokenized_train = train_dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=["text"],
        desc="Tokenizing train dataset",
    )
    
    tokenized_eval = eval_dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=["text"],
        desc="Tokenizing eval dataset",
    )
    
    return tokenized_train, tokenized_eval

# Custom callback to monitor quantization
class QuantizationCallback(TrainerCallback):
    def on_train_begin(self, args, state, control, **kwargs):
        logger.info("Starting Quantization-Aware Training (QAT)")
        
    def on_epoch_begin(self, args, state, control, **kwargs):
        logger.info(f"Starting epoch {state.epoch}")
        
    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step % 100 == 0:
            logger.info(f"Step {state.global_step}: Training with quantized model")



def prepare_for_qat(model):
    logger.info("Preparing model for Quantization-Aware Training")
    
    # Clone the model for quantization
    model_quant = copy.deepcopy(model)
    
    # Cut the model size in half by removing middle layers
    if hasattr(model_quant, "model") and hasattr(model_quant.model, "layers"):
        # For Qwen2 model structure with model.layers
        num_layers = len(model_quant.model.layers)
        keep_layers = num_layers // 2
        
        # Keep first half of the layers
        model_quant.model.layers = model_quant.model.layers[:keep_layers]
        
        logger.info(f"Reduced model from {num_layers} to {keep_layers} layers")
    
    # Create quantizer
    qat_quantizer = Int8DynActInt4WeightQATQuantizer(
        groupsize=groupsize,
        padding_allowed=padding_allowed
    )
    
    # Prepare the model for QAT
    logger.info("Running prepare() to set up QAT layers")
    model_quant = qat_quantizer.prepare(model_quant)
    
    # Rest of your function...
    return model_quant, qat_quantizer



def finalize_quantization(model_quant, qat_quantizer, output_dir):
    print("=============================================", flush=True)
    print("STARTING QUANTIZATION FINALIZATION PROCESS", flush=True)
    print("=============================================", flush=True)
    
    try:
        # Log the models we're working with
        print(f"Model_quant type: {type(model_quant)}", flush=True)
        print(f"QAT Quantizer type: {type(qat_quantizer)}", flush=True)
        
        # Convert to actual quantized operations
        print("ATTEMPTING to convert QAT model to int4...", flush=True)
        model_int4 = qat_quantizer.convert(model_quant)
        print(f"CONVERSION SUCCESSFUL - model_int4 type: {type(model_int4)}", flush=True)
        
        # Create output directories
        int4_output_dir = os.path.join(output_dir, "int4_model")
        print(f"Creating directory: {int4_output_dir}", flush=True)
        os.makedirs(int4_output_dir, exist_ok=True)
        print(f"Directory created successfully: {os.path.exists(int4_output_dir)}", flush=True)
        
        fp32_output_dir = os.path.join(output_dir, "fp32_model")
        print(f"Creating directory: {fp32_output_dir}", flush=True)
        os.makedirs(fp32_output_dir, exist_ok=True)
        print(f"Directory created successfully: {os.path.exists(fp32_output_dir)}", flush=True)
        
        # Save the state dicts with .pt extension
        int4_model_path = os.path.join(int4_output_dir, "model.pt")
        print(f"ATTEMPTING to save int4 quantized model to: {int4_model_path}", flush=True)
        
        # Inspect state dict before saving
        int4_state_dict = model_int4.state_dict()
        print(f"Int4 state dict contains {len(int4_state_dict)} keys", flush=True)
        print(f"First few keys: {list(int4_state_dict.keys())[:3]}", flush=True)
        
        # Save INT4 model
        torch.save(int4_state_dict, int4_model_path)
        print(f"Int4 model SAVED successfully: {os.path.exists(int4_model_path)}", flush=True)
        
        # Save FP32 model
        fp32_model_path = os.path.join(fp32_output_dir, "model.pt")
        print(f"ATTEMPTING to save fp32 model to: {fp32_model_path}", flush=True)
        
        # Inspect original model state dict
        fp32_state_dict = model.state_dict()
        print(f"FP32 state dict contains {len(fp32_state_dict)} keys", flush=True)
        print(f"First few keys: {list(fp32_state_dict.keys())[:3]}", flush=True)
        
        # Save FP32 model
        torch.save(fp32_state_dict, fp32_model_path)
        print(f"FP32 model SAVED successfully: {os.path.exists(fp32_model_path)}", flush=True)
        
        # Verify files exist and compare sizes
        print("CHECKING file sizes...", flush=True)
        if os.path.exists(int4_model_path) and os.path.exists(fp32_model_path):
            # Get file sizes
            fp32_size = os.path.getsize(fp32_model_path) / (1024 ** 2)
            int4_size = os.path.getsize(int4_model_path) / (1024 ** 2)
            
            # Log size comparison
            print(f"FILE SIZE COMPARISON:", flush=True)
            print(f"FP32 Model size: {fp32_size:.2f} MB", flush=True)
            print(f"INT4 Model size: {int4_size:.2f} MB", flush=True)
            
            if fp32_size > 0:  # Avoid division by zero
                reduction = (1 - int4_size/fp32_size) * 100
                print(f"Size reduction: {reduction:.2f}%", flush=True)
            else:
                print("FP32 model size is 0, cannot calculate reduction percentage", flush=True)
        else:
            if not os.path.exists(int4_model_path):
                print(f"INT4 MODEL FILE NOT FOUND at {int4_model_path}", flush=True)
            if not os.path.exists(fp32_model_path):
                print(f"FP32 MODEL FILE NOT FOUND at {fp32_model_path}", flush=True)
            print("Could not compare model sizes; one or both files not found", flush=True)
        
        print("QUANTIZATION FINALIZATION COMPLETED SUCCESSFULLY", flush=True)
        return model_int4
        
    except Exception as e:
        print("!!!! EXCEPTION DURING QUANTIZATION FINALIZATION !!!!", flush=True)
        print(f"Error message: {str(e)}", flush=True)
        print(f"Error type: {type(e).__name__}", flush=True)
        
        # Check what stage we were at when the error occurred
        if 'model_int4' not in locals():
            print("Error occurred during model conversion", flush=True)
        elif not os.path.exists(int4_output_dir):
            print("Error occurred creating output directories", flush=True)
        elif not os.path.exists(int4_model_path):
            print("Error occurred saving int4 model", flush=True)
        elif not os.path.exists(fp32_model_path):
            print("Error occurred saving fp32 model", flush=True)
        
        # Print full traceback
        import traceback
        print("Full traceback:", flush=True)
        traceback.print_exc()
        
        print("QUANTIZATION FINALIZATION FAILED", flush=True)
        return None
                
def main():
    # Prepare dataset
    train_dataset, eval_dataset = prepare_dataset()
    
    # Data collator
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer, 
        mlm=False,  # We're doing causal LM, not masked LM
    )
    
    # Prepare the model for QAT
    model_quant, qat_quantizer = prepare_for_qat(model)


    model_quant.train()

    # Define training arguments - NO DEEPSPEED HERE
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_train_epochs,
        per_device_train_batch_size=per_device_train_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        learning_rate=learning_rate,
        warmup_steps=warmup_steps,
        logging_steps=logging_steps,
        save_steps=save_steps,
        eval_steps=eval_steps,
        save_total_limit=save_total_limit,
        evaluation_strategy="steps",
        load_best_model_at_end=True,
        report_to="wandb",

    )
    
    # Initialize Trainer
    trainer = Trainer(
        model=model_quant,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
        callbacks=[QuantizationCallback()],
    )
    
    # Train the model
    logger.info("Starting QAT training")
    trainer.train()
    
    # Save the trained QAT model
    trainer.save_model(os.path.join(output_dir, "qat_model"))
    
    # Convert the QAT model to int4
    model_int4 = finalize_quantization(model_quant, qat_quantizer, output_dir)
    
    if model_int4:
        logger.info("Quantization completed successfully!")
    else:
        logger.error("Quantization failed!")

if __name__ == "__main__":
    main()
Following training, our final step involves the actual quantization of the model. The finalize_quantization() function converts the QAT-prepared layers into fully quantized modules, embedding int4 grouped per-channel weights and dynamic int8 activation handling into the model.
This function also verifies and compares the sizes of the resulting quantized model against the original FP32 version, highlighting storage reductions and efficiency improvements.

Here are the training logs for my script.

Run: ./qwen_int4_quantized
1


Running inference with the quantized model

Now we will write a script to run inference with the quantized model we fine-tuned previously. The script will load the quantized checkpoint, rebuild the exact quantization configuration used during training, and run inference using a GPU. To correctly reconstruct our INT4 quantized model, the script first initializes the original model architecture, applies the quantizer's prepare() and convert() methods, and finally loads the previously saved INT4 state dictionary.
During inference, we use Hugging Face's generate() method, which efficiently performs sampling-based text generation on GPU, leveraging optimized kernels for INT4 quantized weights and dynamic INT8 activations. We will test the quantized model using several prompts to evaluate its generative capabilities and confirm the inference pipeline is correctly configured.
import torch
import os
import logging
from transformers import AutoTokenizer, AutoModelForCausalLM
from torchao.quantization.prototype.qat import Int8DynActInt4WeightQATQuantizer
import weave; weave.init("QAT")



def load_quantized_model(model_path, model_name="Qwen/Qwen2.5-0.5B"):
    """
    Load the INT4 quantized model from the saved checkpoint
    """
    print(f"Loading quantized model from: {model_path}", flush=True)
    
    # First load the original model architecture
    print(f"Loading base model: {model_name}", flush=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name, 
        torch_dtype=torch.float32
    )
    
    # If you reduced layers during training, do the same reduction here
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        # For Qwen2 model structure with model.layers
        num_layers = len(model.model.layers)
        keep_layers = num_layers // 2
        
        # Keep first half of the layers
        model.model.layers = model.model.layers[:keep_layers]
        print(f"Reduced model from {num_layers} to {keep_layers} layers", flush=True)
    
    # Create a quantizer with the same parameters used during training
    print("Creating quantizer with same parameters as training", flush=True)
    qat_quantizer = Int8DynActInt4WeightQATQuantizer(
        groupsize=224,
        padding_allowed=False
    )
    
    # 1. PREPARE THE MODEL FIRST
    print("1. PREPARING the model with QAT layers", flush=True)
    model = qat_quantizer.prepare(model)
    
    # 2. THEN CONVERT TO INT4 
    print("2. CONVERTING the prepared model to INT4 format", flush=True)
    model_int4 = qat_quantizer.convert(model)
    
    # 3. FINALLY load the state dict
    print(f"3. NOW loading state dict from: {model_path}", flush=True)
    if os.path.exists(model_path):
        state_dict = torch.load(model_path, map_location="cpu")
        print(f"State dict loaded successfully with {len(state_dict)} keys", flush=True)
        
        # Check if state dict keys match model
        model_keys = set(k for k in model_int4.state_dict().keys())
        loaded_keys = set(k for k in state_dict.keys())
        
        # Report on key differences
        missing_keys = model_keys - loaded_keys
        unexpected_keys = loaded_keys - model_keys
        
        if len(missing_keys) > 0:
            print(f"WARNING: {len(missing_keys)} keys are missing in the loaded state dict", flush=True)
            print(f"Sample missing keys: {list(missing_keys)[:5]}", flush=True)
        
        if len(unexpected_keys) > 0:
            print(f"WARNING: {len(unexpected_keys)} unexpected keys in the loaded state dict", flush=True)
            print(f"Sample unexpected keys: {list(unexpected_keys)[:5]}", flush=True)
        
        # Load state dict into model
        model_int4.load_state_dict(state_dict, strict=False)
        print("State dict loaded into model", flush=True)
        
        return model_int4
    else:
        raise FileNotFoundError(f"Model file not found at: {model_path}")


@weave.op 
def run_inference(model, tokenizer, prompt, max_new_tokens=100):
    """
    Run inference with the quantized model
    """
    print(f"\nRunning inference with prompt: '{prompt}'", flush=True)
    
    # Tokenize input
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    
    # Generate text
    print("Generating text...", flush=True)
    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )
    
    # Decode the generated text
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    print("\nGenerated text:", flush=True)
    print("=" * 50, flush=True)
    print(generated_text, flush=True)
    print("=" * 50, flush=True)
    
    return generated_text

def main():
    # Path to the saved quantized model
    int4_model_path = "./qwen_int4_quantized/int4_model/model.pt"
    
    # Original model name
    model_name = "Qwen/Qwen2.5-0.5B"
    
    # Load tokenizer
    print("Loading tokenizer...", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
    
    try:
        # Load and convert the quantized model
        model = load_quantized_model(int4_model_path, model_name)
        
        # Set to evaluation mode
        model.eval()
        
        # Get the best available device dynamically
        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
        print(f"Using device: {device}", flush=True)
        model = model.to(device)
        
        # Run inference with a few test prompts
        prompts = [
            "Once upon a time in a land far away,",
            "The best way to learn machine learning is to",
            "Quantization is important for deploying models because"
        ]
        
        for prompt in prompts:
            generated_text = run_inference(model, tokenizer, prompt, max_new_tokens=100)
            print("\n" + "-" * 30 + "\n", flush=True)
        
        print("Inference completed successfully!", flush=True)
        
    except Exception as e:
        print(f"Error during inference: {str(e)}", flush=True)
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
After running our script, we can navigate to the Weave section within our Weights & Biases project dashboard, where we'll find the generated inference results automatically logged. Because we've decorated our inference function with @weave.op(), each inference output is captured and organized neatly inside Weave.
It's worth noting that this model was trained using only the first half of the pretrained Qwen 2.5 0.5B layers, and on a relatively small number of tokens—so the output here isn’t fully optimized, but reflects what was possible within a constrained training budget.
Training and evaluation logs visualized in Weights & Biases. The model was trained with quantization-aware layers and validated across several test prompts using dynamically quantized activations and INT4 weights.
In Weave, we'll see a clear log of our prompts alongside their generated responses, making it easy to analyze how well our quantized model performed. This allows us to quickly identify any anomalies or artifacts introduced by quantization, ensuring we're confident about the quality of the generated text before deploying the model into a production setting.

Looking ahead: What’s next in LLM quantization?

Quantization is moving fast. Several frameworks and ideas are beginning to shape the future of low-bit training and deployment:
  • BitNet pushes the limits of compression by using ternary weights (–1, 0, +1), achieving just 1.58 bits per parameter on average. It’s a promising path toward ultra-low-resource deployment on edge devices—without requiring specialized hardware.
  • Mixed-precision quantization is becoming more common, allowing more sensitive layers to retain slightly higher bit-widths (e.g., INT6 or INT8), while applying aggressive quantization elsewhere.
  • New numeric formats, like MX4, are gaining traction, especially as hardware like NVIDIA’s upcoming Blackwell GPUs moves away from native INT4 tensor core support.
  • Hybrid approaches are emerging that combine QAT with LoRA or QLoRA. This allows low-rank fine-tuning and quantization to work together—offering a path to highly efficient, domain-adaptable models without compromising quality.
Together, these techniques suggest a future where low-bit models aren’t just efficient, but flexible, trainable, and competitive with their full-precision counterparts.

Conclusion

As models grow larger, efficiency matters more than ever. Without meaningful improvements in efficiency, deploying these powerful models becomes expensive and impractical. Quantization isn't just about reducing memory—it's about making sure bigger models don't become unusable outside of specialized, high-cost environments.
That's why advances in quantization are so important. They're the difference between powerful models being broadly accessible or limited to only the best-equipped organizations. The future of large-scale AI depends on getting efficiency right, ensuring everyone can benefit from these advances, not just those with unlimited resources.
Tags: Articles, Weave, Tutorial, GenAI

Created with ❤️ on Weights & Biases.

https://wandb.ai/byyoung3/Generative-AI/reports/Quantization-Aware-Training-QAT-A-step-by-step-guide-with-PyTorch--VmlldzoxMTk2NTY2Mw
Made with Weights & Biases. Sign up or log in to create reports like this one.
Iterate on AI agents and models faster. Try Weights & Biases today.
Sign upTry W&B now