Implementation:Deepspeedai DeepSpeed Quantization Binding

Knowledge Sources	DeepSpeed
Domains	Quantization, PyTorch_Bindings, Model_Compression
Last Updated	2026-02-09 00:00 GMT

Overview

PyTorch bindings for quantization operations providing 4-bit and 8-bit symmetric/asymmetric quantization with swizzled layouts for distributed training.

Description

This module provides the Python-accessible interface to DeepSpeed's quantization infrastructure through PyBind11. It exposes multiple quantization variants including fake quantization (for QAT), real quantization with parameter extraction, dequantization, and specialized swizzled quantization for communication optimization in distributed training. The swizzled quantization reorders data to optimize all-reduce patterns across pipeline and data parallelism, while LoCo (Low-precision Communication) variants add error feedback for maintaining accuracy over multiple compression steps. The implementation supports both symmetric and asymmetric quantization modes, with configurable bit-widths (4 or 8 bits) and group sizes for balancing compression ratio against accuracy.

Usage

Use these functions from Python to quantize model weights or activations for inference, or to compress gradients for distributed training communication. The swizzled quantization is specifically designed for model parallelism scenarios where gradients need to be reduced across GPUs efficiently.

Code Reference

Source Location

• Repository: DeepSpeed
• File: csrc/quantization/pt_binding.cpp

Signature

// Fake quantization (QAT)
at::Tensor ds_quantize_fp16(at::Tensor& vals, int groups, int bits);
at::Tensor ds_quantize_asym_fp16(at::Tensor& vals, int groups, int bits);

// Real quantization with parameters
std::vector<at::Tensor> quantize_kernel(at::Tensor& input_vals,
                                        int groups, int numBits,
                                        quantize::Type quantType);

// Dequantization
at::Tensor dequantize(at::Tensor& quantized_data, at::Tensor& params,
                     int groups, int num_bits, quantize::Type quant_type);

// Swizzled quantization for distributed training
std::vector<at::Tensor> ds_swizzle_quant(at::Tensor& input_vals,
                                         int groups, int num_bits,
                                         quantize::Type quant_type,
                                         int pipeline_size, int nodes,
                                         int devices_per_node);

// LoCo quantization with error feedback
std::vector<at::Tensor> ds_loco_swizzle_quant(at::Tensor& input_vals,
                                              at::Tensor& error_feedback,
                                              float err_beta, int groups,
                                              int num_bits, quantize::Type quant_type,
                                              int pipeline_size, int nodes,
                                              int devices_per_node);

// Quantized reduction (dequant + reduce + requant)
std::vector<at::Tensor> quantized_reduction(at::Tensor& input_vals,
                                            at::Tensor& input_scales,
                                            int in_groups, int out_groups,
                                            int num_bits, quantize::Type quant_type,
                                            int devices_per_node);

Import

import deepspeed.ops.quantization as ds_quant
from deepspeed.ops.quantization import QuantizationType

I/O Contract

Input	Type	Description
input_vals	torch.Tensor	FP16 tensor to quantize
groups	int	Number of quantization groups
num_bits	int	Bits per element (4 or 8)
quant_type	QuantizationType	Symmetric or Asymmetric mode
error_feedback	torch.Tensor	Error accumulator (LoCo only)

Output	Type	Description
quantized	torch.Tensor (int8)	Packed quantized values
params	torch.Tensor (float)	Quantization parameters (scales, offsets)

Usage Examples

Basic Quantization and Dequantization:

import torch
import deepspeed.ops.quantization as ds_quant
from deepspeed.ops.quantization import QuantizationType

# Quantize weights to 8-bit symmetric
weights = torch.randn(1024, 4096, dtype=torch.float16, device='cuda')
groups = 128  # 32 elements per group (4096/128)

quantized, params = ds_quant.quantize(
    weights, groups=groups, numBits=8,
    quantType=QuantizationType.Symmetric)

print(f"Original size: {weights.element_size() * weights.numel()} bytes")
print(f"Compressed size: {quantized.element_size() * quantized.numel()} bytes")

# Dequantize back
restored = ds_quant.dequantize(
    quantized, params, groups=groups, num_bits=8,
    quant_type=QuantizationType.Symmetric)

error = (weights - restored).abs().mean()
print(f"Reconstruction error: {error.item()}")

4-bit Asymmetric Quantization:

# Higher compression with asymmetric mode
activations = torch.randn(64, 512, 768, dtype=torch.float16, device='cuda')
total_elems = activations.numel()
group_size = 128
groups = total_elems // group_size

# 4-bit asymmetric captures full range better
quantized, params = ds_quant.quantize(
    activations.view(-1), groups=groups, numBits=4,
    quantType=QuantizationType.Asymmetric)

# 8× compression (16-bit → 4-bit = 4×, plus parameters)
compression_ratio = (activations.element_size() * activations.numel()) / \
                   (quantized.element_size() * quantized.numel() +
                    params.element_size() * params.numel())
print(f"Compression ratio: {compression_ratio:.2f}×")

Gradient Compression for Distributed Training:

# Compress gradients before all-reduce
def compress_gradients(param_groups, pipeline_size=2, world_size=8):
    compressed_grads = []

    for param in param_groups:
        if param.grad is None:
            continue

        grad_fp16 = param.grad.half().contiguous()
        num_elems = grad_fp16.numel()
        groups = num_elems // 512  # 512 elements per group

        # Swizzled quantization for optimal all-reduce
        quantized, scales = ds_quant.swizzle_quant(
            grad_fp16, groups=groups, num_bits=8,
            quant_type=QuantizationType.Symmetric,
            pipeline_size=pipeline_size,
            nodes=world_size // 8,  # 8 GPUs per node
            devices_per_node=8)

        compressed_grads.append((quantized, scales, param))

    return compressed_grads

# After all-reduce of compressed data
def decompress_gradients(compressed_grads):
    for quantized, scales, param in compressed_grads:
        grad_fp16 = ds_quant.dequantize(
            quantized, scales,
            groups=scales.size(0), num_bits=8,
            quant_type=QuantizationType.Symmetric)
        param.grad = grad_fp16.float()

LoCo with Error Feedback:

class LocoGradientCompression:
    def __init__(self, model, beta=0.9):
        self.error_feedback = {}
        self.beta = beta
        # Initialize error buffers
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.error_feedback[name] = torch.zeros_like(
                    param.data, dtype=torch.float16)

    def compress(self, name, gradient, groups, num_bits=8):
        # LoCo: gradient + beta * error_feedback
        error_buf = self.error_feedback[name]

        quantized, scales = ds_quant.loco_swizzle_quant(
            gradient.half().contiguous(),
            error_buf, self.beta, groups, num_bits,
            QuantizationType.Symmetric,
            pipeline_size=2, nodes=1, devices_per_node=8)

        return quantized, scales

# Usage
compressor = LocoGradientCompression(model, beta=0.9)
for name, param in model.named_parameters():
    if param.grad is not None:
        groups = param.grad.numel() // 256
        q_grad, scales = compressor.compress(name, param.grad, groups)
        # ... all-reduce q_grad and scales ...

Fake Quantization for QAT:

# Quantization-Aware Training (simulates quantization noise)
class FakeQuantize(torch.nn.Module):
    def __init__(self, num_bits=8, groups=128):
        super().__init__()
        self.num_bits = num_bits
        self.groups = groups

    def forward(self, x):
        if self.training:
            # Simulate quantization during forward pass
            x_quant = ds_quant.ds_quantize_fp16(
                x.contiguous(), self.groups, self.num_bits)
            return x_quant
        else:
            return x

# Apply to activations in model
model.intermediate_act = torch.nn.Sequential(
    model.intermediate_act,
    FakeQuantize(num_bits=8, groups=64)
)

Related Pages

• Quantization Utils - CUDA implementation details
• Dequantization Utils - Reverse operation
• FP Quantize - Floating-point format quantization