xbe.at

LLMs and generative AI / Vector databases and embeddings

Quantization for Efficient Large Language Models

10 min read, 15 sections

Contents
  1. Introduction to Quantization in LLMs
  2. Types of Quantization
  3. Benefits of Quantization in LLMs
  4. Implementing Post-Training Quantization
  5. Quantization-Aware Training
  6. Quantization Granularity
  7. Handling Activation Functions
  8. Quantization for Inference
  9. Quantization-Friendly Model Design
  10. Quantization for Transformers
  11. Quantization and Fine-tuning
  12. Real-Life Example: Text Classification
  13. Real-Life Example: Image Classification
  14. Challenges and Limitations of Quantization
  15. Additional Resources

Introduction to Quantization in LLMs

Quantization is a technique that reduces the precision of model weights and activations, allowing for more efficient storage and computation of Large Language Models. This process is crucial for deploying LLMs on resource-constrained devices and improving inference speed without significant loss in model performance.

import torch

# Original float32 tensor
original_tensor = torch.randn(1000, 1000)

# Quantize to int8
quantized_tensor = torch.quantize_per_tensor(original_tensor, scale=0.1, zero_point=0, dtype=torch.qint8)

print(f"Original size: {original_tensor.element_size() * original_tensor.nelement() / 1024:.2f} KB")
print(f"Quantized size: {quantized_tensor.element_size() * quantized_tensor.nelement() / 1024:.2f} KB")

Types of Quantization

There are several types of quantization techniques, including post-training quantization (PTQ) and quantization-aware training (QAT). PTQ is applied after model training, while QAT incorporates quantization during the training process for better accuracy.

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)
    
    def forward(self, x):
        return self.fc(x)

# Create and quantize model
model = SimpleModel()
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

print(f"Original model size: {sum(p.numel() for p in model.parameters()) * 4 / 1024:.2f} KB")
print(f"Quantized model size: {sum(p.numel() for p in quantized_model.parameters()) * 1 / 1024:.2f} KB")

Benefits of Quantization in LLMs

Quantization offers several advantages for LLMs, including reduced memory footprint, faster inference times, and lower power consumption. These benefits make it possible to deploy large models on edge devices and improve overall system efficiency.

import time
import torch
import torch.nn as nn

class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(*[nn.Linear(1000, 1000) for _ in range(10)])
    
    def forward(self, x):
        return self.fc(x)

model = LargeModel()
input_tensor = torch.randn(100, 1000)

# Measure inference time for original model
start_time = time.time()
_ = model(input_tensor)
original_time = time.time() - start_time

# Quantize model
quantized_model = torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)

# Measure inference time for quantized model
start_time = time.time()
_ = quantized_model(input_tensor)
quantized_time = time.time() - start_time

print(f"Original inference time: {original_time:.4f} seconds")
print(f"Quantized inference time: {quantized_time:.4f} seconds")
print(f"Speedup: {original_time / quantized_time:.2f}x")

Implementing Post-Training Quantization

Post-Training Quantization (PTQ) is a straightforward method to quantize pre-trained models. It involves converting the model's weights and activations to lower precision without retraining.

import torch
import torch.nn as nn

class ExampleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(64 * 6 * 6, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Create and quantize model
model = ExampleModel()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)
torch.quantization.convert(model, inplace=True)

print(f"Quantized model: {model}")

Quantization-Aware Training

Quantization-Aware Training (QAT) simulates the effects of quantization during the training process, allowing the model to adapt to the reduced precision and potentially achieve better accuracy compared to post-training quantization.

import torch
import torch.nn as nn
import torch.optim as optim

class QAwareModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(64 * 6 * 6, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Prepare model for QAT
model = QAwareModel()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)

# Training loop (simplified)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

for epoch in range(10):
    # ... training code ...
    pass

# Convert to quantized model
torch.quantization.convert(model, inplace=True)

print(f"Quantization-aware trained model: {model}")

Quantization Granularity

Quantization can be applied at different levels of granularity, such as per-tensor, per-channel, or even per-group. Finer granularity often leads to better accuracy but may increase complexity.

import torch
import torch.nn as nn

# Create a sample tensor
tensor = torch.randn(2, 3, 4, 4)

# Per-tensor quantization
per_tensor_scale, per_tensor_zero_point = torch.quantization.calculate_qparams(tensor)
per_tensor_quantized = torch.quantize_per_tensor(tensor, per_tensor_scale, per_tensor_zero_point, torch.quint8)

# Per-channel quantization
per_channel_scales, per_channel_zero_points = torch.quantization.calculate_qparams(tensor, axis=1)
per_channel_quantized = torch.quantize_per_channel(tensor, per_channel_scales, per_channel_zero_points, 1, torch.quint8)

print(f"Per-tensor quantized shape: {per_tensor_quantized.shape}")
print(f"Per-channel quantized shape: {per_channel_quantized.shape}")

Handling Activation Functions

Proper handling of activation functions is crucial in quantized models. Some activation functions, like ReLU, can be easily quantized, while others may require special treatment.

import torch
import torch.nn as nn

class QuantizedReLU(nn.Module):
    def __init__(self):
        super().__init__()
        self.activation_post_process = torch.quantization.default_dynamic_quant_observer()

    def forward(self, x):
        x = torch.relu(x)
        x = self.activation_post_process(x)
        return x

# Create a quantized model with custom ReLU
class CustomQuantizedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.relu = QuantizedReLU()
        self.fc = nn.Linear(64 * 6 * 6, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = CustomQuantizedModel()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)
torch.quantization.convert(model, inplace=True)

print(f"Custom quantized model: {model}")

Quantization for Inference

Quantization is particularly useful for inference, where model size and speed are critical. Here's an example of how to use a quantized model for inference.

import torch
import torchvision.models as models

# Load a pre-trained ResNet model
model = models.resnet18(pretrained=True)

# Quantize the model
model.eval()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)
torch.quantization.convert(model, inplace=True)

# Prepare input
input_tensor = torch.randn(1, 3, 224, 224)

# Perform inference
with torch.no_grad():
    output = model(input_tensor)

print(f"Output shape: {output.shape}")
print(f"Top 5 predictions: {torch.topk(output, 5).indices}")

Quantization-Friendly Model Design

Designing models with quantization in mind can lead to better performance after quantization. This includes choosing appropriate activation functions and avoiding operations that are difficult to quantize.

import torch
import torch.nn as nn

class QuantizationFriendlyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.adaptive_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.adaptive_avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = QuantizationFriendlyModel()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)
torch.quantization.convert(model, inplace=True)

print(f"Quantization-friendly model: {model}")

Quantization for Transformers

Transformers, the backbone of many LLMs, can benefit significantly from quantization. Here's an example of quantizing a simple transformer model.

import torch
import torch.nn as nn

class SimpleTransformer(nn.Module):
    def __init__(self, d_model=512, nhead=8):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead)
        self.norm1 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, src):
        src2 = self.self_attn(src, src, src)[0]
        src = src + src2
        src = self.norm1(src)
        src2 = self.ff(src)
        src = src + src2
        src = self.norm2(src)
        return src

model = SimpleTransformer()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)
torch.quantization.convert(model, inplace=True)

print(f"Quantized Transformer model: {model}")

Quantization and Fine-tuning

When quantizing pre-trained models, it's often beneficial to fine-tune the quantized model on a small dataset to recover any lost accuracy.

import torch
import torch.nn as nn
import torch.optim as optim

# Assume we have a pre-trained quantized model
quantized_model = torch.load('path_to_quantized_model.pth')

# Prepare data loader
train_loader = torch.utils.data.DataLoader(...)  # Your dataset here

# Fine-tuning loop
optimizer = optim.Adam(quantized_model.parameters(), lr=0.0001)
criterion = nn.CrossEntropyLoss()

for epoch in range(5):  # Fine-tune for 5 epochs
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = quantized_model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}')

# Save the fine-tuned quantized model
torch.save(quantized_model, 'fine_tuned_quantized_model.pth')

Real-Life Example: Text Classification

Let's implement a quantized LSTM model for text classification, demonstrating how quantization can be applied to natural language processing tasks.

import torch
import torch.nn as nn

class QuantizedLSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, text):
        embedded = self.embedding(text)
        output, (hidden, cell) = self.lstm(embedded)
        hidden = hidden[-1, :, :]
        return self.fc(hidden)

# Create model
model = QuantizedLSTMClassifier(vocab_size=10000, embedding_dim=100, hidden_dim=256, output_dim=2)

# Prepare for quantization
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)

# Training loop would go here

# Convert to fully quantized model
torch.quantization.convert(model, inplace=True)

# Example usage
input_text = torch.randint(0, 10000, (32, 100))  # Batch of 32 sentences, each with 100 words
output = model(input_text)
print(f"Output shape: {output.shape}")

Real-Life Example: Image Classification

Let's implement a quantized CNN model for image classification, demonstrating how quantization can be applied to computer vision tasks.

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
from torch.utils.data import DataLoader

class QuantizedCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.fc = nn.Linear(64 * 8 * 8, 10)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = self.fc(x)
        return x

# Create and quantize model
model = QuantizedCNN()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)

# Load CIFAR10 dataset
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=32, shuffle=True)

# Training loop (simplified)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(5):
    for inputs, labels in trainloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# Convert to fully quantized model
torch.quantization.convert(model, inplace=True)

# Test the model
testset = CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = DataLoader(testset, batch_size=32, shuffle=False)

correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in testloader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy on test set: {100 * correct / total}%')

Challenges and Limitations of Quantization

While quantization offers significant benefits, it also comes with challenges. These include potential accuracy loss, especially for complex models or tasks, and the need for careful tuning of quantization parameters.

import torch
import torch.nn as nn

class ComplexModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(64, 128, 3)
        self.bn2 = nn.BatchNorm2d(128)
        self.fc1 = nn.Linear(128 * 5 * 5, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool(self.relu(self.bn1(self.conv1(x))))
        x = self.pool(self.relu(self.bn2(self.conv2(x))))
        x = x.view(-1, 128 * 5 * 5)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = ComplexModel()

# Attempt quantization
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)

# This may result in warnings or errors due to unsupported operations
try:
    torch.quantization.convert(model, inplace=True)
except Exception as e:
    print(f"Quantization failed: {e}")
    print("Complex models may require careful handling of each layer for successful quantization.")

Additional Resources

For further exploration of quantization techniques and their application to Large Language Models, consider the following resources:

  1. "Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference" by Jacob et al. (2018) ArXiv: https://arxiv.org/abs/1712.05877
  2. "Quantizing deep convolutional networks for efficient inference: A whitepaper" by Krishnamoorthi (2018) ArXiv: https://arxiv.org/abs/1806.08342
  3. "ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers" by Yao et al. (2022) ArXiv: https://arxiv.org/abs/2206.01861

These papers provide in-depth discussions on quantization techniques, their implementation, and their impact on model performance.

More in Vector databases and embeddings

  1. Running Open-Source LLMs Locally with Ollama10 min

    The first step involves configuring the Python environment to interact with Ollama's API endpoints. This allows seamless communication between Python...

  2. Optimizing LLM Inference Costs with Mixed GPU Types9 min

    Mixed GPU inference is a technique to optimize LLM (Large Language Model) inference costs by utilizing different types of GPUs. This approach leverages the...

  3. Techniques for Compressing Large Language Models in Python7 min

    Model compression is crucial for deploying large language models (LLMs) in resource-constrained environments. It aims to reduce model size and computational...

  4. Open-Source Tools for Scaling Testing and Deploying LLMs6 min

    The development and deployment of Large Language Models (LLMs) require a robust set of tools for scaling, testing, deployment, and monitoring. This...

  5. Tokens and Tokenization in Large Language Models in Python7 min

    Introduction to Tokens and Tokenization in LLMs

  6. MatMul-Free LLMs for Enhanced Efficiency in Python11 min

    Traditional Large Language Models (LLMs) rely heavily on matrix multiplication operations, which can be computationally expensive. MatMul-Free LLMs aim to...