Implementation:Deepspeedai DeepSpeed FP Quantize

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

Overview

PyTorch bindings for floating-point format quantization supporting FP4, FP6, FP8, and FP12 with configurable mantissa bits and stochastic rounding.

Description

This module provides Python-accessible quantization to custom floating-point formats with reduced precision compared to FP16/FP32. Unlike integer quantization, floating-point quantization maintains the dynamic range properties of floating-point representation but with fewer mantissa and exponent bits. The implementation supports four custom formats: FP4 (1 sign + 2 exponent + 1 mantissa), FP6 (1+2+3), FP8 (1+4+3 or 1+5+2), and FP12 (1+4+7), each with different precision/range tradeoffs. The quantization operates on groups of elements to amortize the overhead of storing quantization metadata. Stochastic rounding is optionally available to reduce quantization bias during iterative processes like training. The dequantization function reconstructs FP16/BF16 values from the compressed format, and selective dequantization allows efficiently extracting specific elements without decompressing the entire tensor.

Usage

Use these functions when implementing custom floating-point precision training or inference, particularly useful for research into ultra-low precision neural networks. FP8 is becoming popular for training acceleration, while FP6 and FP4 are experimental formats for extreme compression.

Code Reference

Source Location

Repository: DeepSpeed
File: csrc/fp_quantizer/fp_quantize.cpp

Signature

// Quantize to custom FP format
at::Tensor quantize(torch::Tensor& out, torch::Tensor& val,
                   int group_size, int stochastic_rounding,
                   int q_bits, int q_mantisa_bits);

// Dequantize from custom FP format
void dequantize(torch::Tensor& val, torch::Tensor& val_q,
               int group_size, int q_mantisa_bits, int q_exponent_bits);

// Selective dequantization (sparse access)
void selective_dequantize(torch::Tensor& val, torch::Tensor& val_q,
                         torch::Tensor& indexes, int group_size,
                         int q_mantisa_bits, int q_exponent_bits);

// Extract scale factors
at::Tensor get_scales(torch::Tensor& out, int num_groups);

Import

import deepspeed.ops.fp_quantizer as fp_quant

I/O Contract

Input	Type	Description
val	torch.Tensor	FP16/BF16 tensor to quantize
group_size	int	Elements per quantization group
q_bits	int	Total bits: 4, 6, 8, or 12
q_mantisa_bits	int	Mantissa bits (1-7 depending on format)
stochastic_rounding	int	0=nearest, 1=stochastic

Output	Type	Description
out	torch.Tensor (uint8)	Quantized values + scales
scales	torch.Tensor (float)	Per-group scale factors

Usage Examples

FP8 Quantization (E4M3):

import torch
import deepspeed.ops.fp_quantizer as fp_quant

# FP8 E4M3: 1 sign + 4 exponent + 3 mantissa
weights = torch.randn(1024, 4096, dtype=torch.float16, device='cuda')
group_size = 128

# Allocate output buffer (includes scales)
total_elems = weights.numel()
num_groups = total_elems // group_size
# Each element: 8 bits, plus 4 bytes per group for scale
output_size = total_elems + num_groups * 4
out = torch.empty(output_size, dtype=torch.uint8, device='cuda')

# Quantize to FP8
fp_quant.quantize(
    out, weights.view(-1),
    group_size=group_size,
    stochastic_rounding=0,
    q_bits=8,
    q_mantisa_bits=3  # E4M3 format
)

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

# Dequantize
restored = torch.empty_like(weights)
fp_quant.dequantize(
    restored.view(-1), out,
    group_size=group_size,
    q_mantisa_bits=3,
    q_exponent_bits=4
)

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

FP6 for Extreme Compression:

# FP6: 1 sign + 2 exponent + 3 mantissa
# Range: ~[0, 28] with coarse precision
activations = torch.randn(64, 512, 768, dtype=torch.float16, device='cuda')

# Flatten and quantize
flat_acts = activations.view(-1)
group_size = 256
num_groups = flat_acts.numel() // group_size

# FP6 packs into 6 bits per value
bits_per_group = group_size * 6
bytes_per_group = (bits_per_group + 7) // 8
output_size = num_groups * (bytes_per_group + 4)  # +4 for scale

out = torch.empty(output_size, dtype=torch.uint8, device='cuda')

fp_quant.quantize(
    out, flat_acts,
    group_size=group_size,
    stochastic_rounding=0,
    q_bits=6,
    q_mantisa_bits=3
)

# ~4× compression (16-bit → 6-bit effective)
compression = (flat_acts.element_size() * flat_acts.numel()) / \
             (out.element_size() * out.numel())
print(f"Compression ratio: {compression:.2f}×")

Stochastic Rounding for Training:

# Stochastic rounding reduces bias in iterative updates
class FP8Optimizer:
    def __init__(self, params, lr=1e-3):
        self.params = list(params)
        self.lr = lr

    def step(self):
        for param in self.params:
            if param.grad is None:
                continue

            # Compute update in FP16
            update = param.grad * self.lr

            # Quantize update with stochastic rounding
            group_size = 128
            num_groups = update.numel() // group_size
            out_size = update.numel() + num_groups * 4

            quantized = torch.empty(out_size, dtype=torch.uint8,
                                   device=update.device)

            fp_quant.quantize(
                quantized, update.view(-1),
                group_size=group_size,
                stochastic_rounding=1,  # Enable stochastic
                q_bits=8, q_mantisa_bits=3
            )

            # Dequantize and apply
            dequant = torch.empty_like(update)
            fp_quant.dequantize(dequant.view(-1), quantized,
                               group_size, 3, 4)

            param.data -= dequant
            param.grad.zero_()

Selective Dequantization for Sparse Access:

# Only dequantize specific elements (e.g., for sampling)
full_tensor = torch.randn(100000, dtype=torch.float16, device='cuda')

# Quantize full tensor
group_size = 512
num_groups = full_tensor.numel() // group_size
out = torch.empty(full_tensor.numel() + num_groups * 4,
                 dtype=torch.uint8, device='cuda')

fp_quant.quantize(out, full_tensor, group_size, 0, 8, 3)

# Later: access only specific indices
sample_indices = torch.randint(0, full_tensor.numel(), (1000,),
                              device='cuda', dtype=torch.int32)

# Selective dequantization (more efficient than full decompression)
sampled_values = torch.empty(1000, dtype=torch.float16, device='cuda')
fp_quant.selective_dequantize(
    sampled_values, out, sample_indices, group_size, 3, 4
)

# Verify
full_dequant = torch.empty_like(full_tensor)
fp_quant.dequantize(full_dequant, out, group_size, 3, 4)
print("Match:", torch.allclose(sampled_values,
                                full_dequant[sample_indices]))

Extract Scales for Analysis:

# Quantize with scale extraction
weights = torch.randn(2048, 8192, dtype=torch.float16, device='cuda')
group_size = 256
num_groups = weights.numel() // group_size

out = torch.empty(weights.numel() + num_groups * 4,
                 dtype=torch.uint8, device='cuda')
fp_quant.quantize(out, weights.view(-1), group_size, 0, 8, 3)

# Extract scale factors for each group
scales = fp_quant.get_scales(out, num_groups)

# Analyze distribution
print(f"Scale stats:")
print(f"  Min: {scales.min().item():.6f}")
print(f"  Max: {scales.max().item():.6f}")
print(f"  Mean: {scales.mean().item():.6f}")
print(f"  Std: {scales.std().item():.6f}")

# Large scale variation indicates heterogeneous ranges

FP4 Ultra-Compression:

# FP4: 1 sign + 2 exponent + 1 mantissa (very limited range/precision)
# Best for normalized data with small dynamic range

# Normalize to [-1, 1] range
data = torch.tanh(torch.randn(1024, 1024, dtype=torch.float16, device='cuda'))

group_size = 64
num_groups = data.numel() // group_size

# FP4 packs 4 bits per value = 0.5 bytes
bits_total = data.numel() * 4 + num_groups * 32  # +32 bits per scale
bytes_total = (bits_total + 7) // 8

out = torch.empty(bytes_total, dtype=torch.uint8, device='cuda')

fp_quant.quantize(out, data.view(-1), group_size, 0,
                 q_bits=4, q_mantisa_bits=1)

# 8× compression!
compression = (data.element_size() * data.numel()) / bytes_total
print(f"FP4 compression: {compression:.1f}×")

Related Pages

Quantization Binding - Integer quantization variant
Quantization Utils - Device-side quantization
Conversion Utils - Type conversion utilities

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment