Low-Bit Precision Training in PyTorch: Techniques and Code Examples

In the era of large-scale deep learning, efficient training is critical. Low-bit precision training — using fewer bits to represent weights, activations, and gradients — can dramatically speed up computation, reduce memory usage, and enable deployment to resource-constrained devices, all while preserving most model accuracy. In PyTorch, multiple methods for low-bit precision exist: built-in quantization workflows, mixed precision, and community libraries.

Let’s walk through the most popular techniques, with clear explanations and ready-to-run code samples.

#1. Quantization in PyTorch

Quantization refers to converting high-precision numbers (often FP32) into lower-precision (e.g., INT8, INT4). PyTorch supports several workflows:

#1.1 Dynamic Quantization

Dynamic quantization is simple and works best for models with large linear layers (e.g., transformers, LSTMs). The model weights are quantized post-training, while activations are quantized on-the-fly during inference.

Steps:

1. Build or load a model.
2. Apply torch.quantization.quantize_dynamic().
3. Use quantized model for inference.

Example:

import torch
import torch.nn as nn
from torch.quantization import quantize_dynamic
# Build a simple LSTM-based model
model = nn.LSTM(input_size=10, hidden_size=20, num_layers=2)
# Quantize model to INT8
quantized_model = quantize_dynamic(
    model,  # Model instance
    {nn.LSTM, nn.Linear},  # Layers to quantize
    dtype=torch.qint8
)
# Use as normal for inference
data = torch.randn(5, 3, 10)  # seq_len, batch, input_size
output, _ = quantized_model(data)
print(output.shape)  # (5, 3, 20)

Pros: Zero code changes to your training pipeline
Cons: Only for inference; limited control over which layers are quantized.

#1.2 Static Quantization (Post-Training Quantization)

Static quantization quantizes both weights and activations. It requires model calibration with sample data to collect quantization parameters.

Steps:

1. Fuse layers (Conv+BN+ReLU for example)
2. Prepare the model
3. Calibrate the model with representative data
4. Convert to a quantized model

Example:

import torch
import torch.nn as nn

Pros: Often higher accuracy than dynamic quantization
Cons: Calibration required; works mainly for inference

#1.3 Quantization-Aware Training (QAT)

QAT simulates quantization during training, so the model can learn to compensate for quantization effects. This approach gives the best accuracy after quantization.

Steps:

1. Prepare your model for QAT with prepare_qat()
2. Train your model (as usual, for a few epochs)
3. Convert to a quantized model for deployment

Example:

import torch
from torch.quantization import prepare_qat, convert

Pros: Best accuracy; robust to quantization noise
Cons: Requires more effort; model needs to be built with QAT support

#2. Mixed Precision Training (FP16/BF16)

“Mixed precision” relies on half-precision floating point arithmetic (FP16 or bfloat16) to speed up model training using GPU tensor cores. This is not integer quantization, but is often classified as “low-precision training”.

Steps:

1. Use torch.cuda.amp for automatic mixed precision (AMP)
2. Train as usual; PyTorch automatically handles precision casting

Example:

import torch
import torch.nn as nn

Pros: Trains faster on supported hardware (NVIDIA GPUs); minimal code changes
Cons: Less memory savings than int8; does not work on CPUs

#3. Low-Bit Libraries: 8-Bit and 4-Bit Training/Inference

Some PyTorch extension libraries support even more aggressive quantization, down to 8 and 4 bits, and enable efficient training.

#a) BitsAndBytes (bnb) / HuggingFace Accelerate

BitsAndBytes supports 8-bit and 4-bit quantized training for Transformers and other models.

Example:

# Install: pip install bitsandbytes

Pros: Ultra-low memory footprint, good for deploying LLMs
Cons: Limited library support; requires special hardware/software

#4. Custom Quantization / Research Code

For 4-bit or lower training, you typically need to modify models or gradients, or use research libraries like Ladder or implement custom quantization.

A minimal custom 4-bit quantization example:

def quantize_to_4bit(tensor):
    # Fake quantization example: scales tensor to [0, 15], rounds, rescales

Usually, these methods require custom layers and are for advanced use-cases.

#Summary Table

#Conclusion

Low-bit precision training is an essential tool for scaling deep learning while keeping resource usage in check. With PyTorch, you can take advantage of simple quantization workflows for deployment, powerful quantization-aware training for accuracy, and use advanced libraries for extreme memory savings.

Originally published on Medium.