【DeepSpeed】模型压缩: 层减少、权重量化、激活量化、剪枝(稀疏剪枝、行剪枝、头剪枝、通道剪枝)以及 ZeroQuant 和 XTC
DeepSpeed 提供了一个强大的模型压缩库(DeepSpeed Compression),旨在通过多种压缩技术降低深度学习模型的内存占用、推理延迟和计算需求,同时尽量保持模型性能。该库特别适合大规模 Transformer 模型(如 LLaMA、BERT、GPT),支持分布式训练和推理环境。DeepSpeed 的压缩技术包括 层减少、权重量化、激活量化、剪枝(稀疏剪枝、行剪枝、头剪枝、通道剪枝
DeepSpeed 提供了一个强大的模型压缩库(DeepSpeed Compression),旨在通过多种压缩技术降低深度学习模型的内存占用、推理延迟和计算需求,同时尽量保持模型性能。该库特别适合大规模 Transformer 模型(如 LLaMA、BERT、GPT),支持分布式训练和推理环境。DeepSpeed 的压缩技术包括 层减少、权重量化、激活量化、剪枝(稀疏剪枝、行剪枝、头剪枝、通道剪枝)以及专用框架如 ZeroQuant 和 XTC。
1. DeepSpeed 模型压缩库概述
1.1 功能
DeepSpeed Compression Library 旨在通过以下方式优化模型:
- 模型体积压缩:减少模型参数存储需求(如 32x 体积压缩,精度损失极小)。
- 推理加速:降低推理延迟(如 2x 延迟减少),提高吞吐量。
- 压缩成本降低:如 ZeroQuant 技术将压缩成本降低 5000x。
- 端到端优化:通过高效推理引擎(如 DeepSpeed-Inference)提升压缩模型的计算效率。
- 可组合性:支持多种压缩方法的协同应用(如层减少和量化),通过统一的 API 简化操作。
- 易用性:提供模块化设计,开发者可轻松添加自定义压缩方法。
1.2 适用场景
- 大模型部署:如 LLaMA-7B、Mixtral-8x7B,在边缘设备或资源受限环境部署。
- 实时推理:如对话系统、文本生成,需要低延迟和高吞吐量。
- 成本敏感场景:如云服务,需降低存储和计算成本。
- 研究与开发:探索新型压缩技术或优化现有模型。
1.3 支持的压缩方法
DeepSpeed Compression 目前支持以下七种压缩方法:
- 层减少(Layer Reduction):通过知识蒸馏减少隐藏层数量,保持网络宽度。
- 权重量化(Weight Quantization):将权重从 FP16/FP32 量化为 INT8/INT4,甚至二进制/三进制。
- 激活量化(Activation Quantization):将激活值量化为 INT8,减少内存和计算开销。
- 稀疏剪枝(Sparse Pruning):移除不重要的权重,生成稀疏矩阵。
- 行剪枝(Row Pruning):删除权重矩阵的整行。
- 头剪枝(Head Pruning):移除 Transformer 模型中不重要的注意力头。
- 通道剪枝(Channel Pruning):减少卷积层或线性层的通道数。
1.4 关键技术
- ZeroQuant:零成本量化技术,支持 INT4/INT8 混合精度量化,显著降低压缩时间(如 GPT-3-1.3B 在单 GPU 上 3 小时完成量化,成本降低 5000x)。
- XTC(Extreme Compression):结合层减少、二值化和知识蒸馏,实现 32x 模型体积压缩,精度损失最小。
1.5 压缩组合器
通过简单 API 自动组合多种压缩方法,简化开发流程。
from deepspeed.compression import init_compression, redundancy_clean
# 初始化压缩
model = init_compression(model, ds_config)
# 训练后清理压缩
model = redundancy_clean(model)
2. 通用教程
DeepSpeed 压缩库通过 compression composer 提供统一 API,简化压缩方法组合。需安装 DeepSpeed(版本 >= 0.7.0):
pip install deepspeed>=0.7.0
2.1 层减少(Layer Reduction)
说明:层减少通过知识蒸馏减少 Transformer 模型层数(如将 BERT-base 从 12 层减少到 5 层),降低计算复杂度和模型大小,尽量保留性能。
JSON 配置(deepspeed_config.json):
{
"compression_training": {
"layer_reduction": {
"enabled": true, // 是否启用层减少,布尔值,true 表示启用
"keep_number_of_layers": 5, // 目标保留的层数,整数,需小于原始层数(如 BERT-base 为 12)
"teacher_layer_mapping": [0, 2, 4, 7, 11] // 教师模型层到学生模型层的映射,列表,指定哪些层保留
}
}
}
JSON 配置解释:
enabled: 控制是否启用层减少功能。设为true启用,false禁用。keep_number_of_layers: 指定压缩后保留的层数,必须为正整数且小于原始模型层数。例如,BERT-base 有 12 层,设为 5 表示压缩到 5 层。teacher_layer_mapping: 定义教师模型(原始模型)到学生模型(压缩模型)的层映射,列表中的整数表示教师模型的层索引。例如,[0, 2, 4, 7, 11]表示从 12 层中选择第 0、2、4、7、11 层映射到学生模型的 5 层,用于知识蒸馏。
示例代码:
import deepspeed
from transformers import BertModel
# 加载预训练模型
model = BertModel.from_pretrained("bert-base-uncased")
# 初始化 DeepSpeed 压缩
deepspeed_config = "deepspeed_config.json"
model = deepspeed.compression.init_compression(model, deepspeed_config)
# 执行层减少
model = deepspeed.compression.redundancy_clean(model, deepspeed_config)
# 模型推理
input_ids = ... # 输入数据
outputs = model(input_ids)
2.2 权重量化(Weight Quantization)
说明:权重量化将权重从高精度(如 FP32)转换为低精度(如 INT8、INT4),减少模型大小和推理时间。支持对称和非对称量化。
JSON 配置:
{
"compression_training": {
"weight_quantization": {
"shared_parameters": {
"enabled": true, // 是否启用权重量化,布尔值,true 表示启用
"quantizer_kernel": false, // 是否使用量化内核加速,布尔值,true 表示启用(需硬件支持)
"quantization_type": "asymmetric", // 量化类型,字符串,可选 "symmetric" 或 "asymmetric"
"start_bits": 12, // 初始量化位宽,整数,范围 [1, 16]
"target_bits": 8, // 目标量化位宽,整数,范围 [1, 16],需小于等于 start_bits
"quantization_period": 50 // 量化周期,整数,表示每多少步执行一次量化
},
"different_groups": {
"wq1": {
"params": {
"start_bits": 12, // 组内初始量化位宽
"target_bits": 8, // 组内目标量化位宽
"quantization_period": 50 // 组内量化周期
},
"modules": ["encoder.layer.0", "encoder.layer.1"] // 应用量化的模块名列表
}
}
}
}
}
JSON 配置解释:
shared_parameters: 全局量化参数,应用于所有未单独配置的模块。enabled: 控制是否启用权重量化。quantizer_kernel: 是否使用硬件加速量化内核(需特定硬件支持,如 NVIDIA GPU)。quantization_type: 量化类型,symmetric表示对称量化(权重范围对称于零),asymmetric表示非对称量化(权重范围动态调整)。start_bits: 初始量化位宽,通常为较高精度(如 12 位),用于渐进量化。target_bits: 目标量化位宽,如 8 位或 4 位,决定最终模型大小。quantization_period: 量化更新的步数间隔,控制量化过程的渐进性。
different_groups: 针对特定模块的差异化量化配置。wq1: 量化组名称,用户自定义。params: 组内量化参数,与shared_parameters字段相同。modules: 指定应用该组量化的模块名称(如 Transformer 的特定层)。
示例代码:
import deepspeed
import torch
from transformers import BertModel
# 加载模型
model = BertModel.from_pretrained("bert-base-uncased").to("cuda")
model.eval()
# DeepSpeed 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = deepspeed.compression.init_compression(model, deepspeed_config)
# 执行权重量化
model = deepspeed.compression.redundancy_clean(model, deepspeed_config)
# 推理
data = torch.rand(4, 128, device="cuda")
outputs = model(data)
2.3 激活量化(Activation Quantization)
说明:激活量化将中间层激活值量化为低精度(如 INT8),减少计算开销。需校准数据以确定量化范围。
JSON 配置:
{
"compression_training": {
"activation_quantization": {
"shared_parameters": {
"enabled": true, // 是否启用激活量化,布尔值,true 表示启用
"quantization_type": "symmetric", // 量化类型,字符串,可选 "symmetric" 或 "asymmetric"
"start_bits": 16, // 初始量化位宽,整数,范围 [1, 16]
"target_bits": 8, // 目标量化位宽,整数,范围 [1, 16]
"quantization_period": 100 // 量化周期,整数,表示每多少步执行一次量化
}
}
}
}
JSON 配置解释:
shared_parameters:enabled: 控制是否启用激活量化。quantization_type: 量化类型,symmetric表示对称量化,asymmetric表示非对称量化。start_bits: 初始激活量化位宽,通常为 16 位,确保初始精度。target_bits: 目标激活量化位宽,如 8 位,决定计算效率。quantization_period: 量化更新间隔,控制量化渐进性。
示例代码:
import deepspeed
import torch
from transformers import BertModel
# 加载模型
model = BertModel.from_pretrained("bert-base-uncased").to("cuda")
model.eval()
# DeepSpeed 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = deepspeed.compression.init_compression(model, deepspeed_config)
# 执行激活量化
model = deepspeed.compression.redundancy_clean(model, deepspeed_config)
# 推理
data = torch.rand(4, 128, device="cuda")
outputs = model(data)
2.4 剪枝(Pruning)
剪枝移除模型中不重要的权重或结构,减少模型大小和计算量。DeepSpeed 支持以下四种剪枝方式:
2.4.1 稀疏剪枝(Sparse Pruning)
说明:稀疏剪枝移除权重矩阵中的个别元素,生成稀疏矩阵。适合通用场景,但硬件加速效果有限。
JSON 配置:
{
"compression_training": {
"sparse_pruning": {
"enabled": true, // 是否启用稀疏剪枝,布尔值,true 表示启用
"sparse_ratio": 0.5, // 稀疏比率,浮点数,范围 [0, 1],表示移除的权重比例
"schedule_offset": 0 // 剪枝开始的步数偏移,整数,控制剪枝时机
}
}
}
JSON 配置解释:
enabled: 控制是否启用稀疏剪枝。sparse_ratio: 移除权重的比例,如 0.5 表示移除 50% 的权重。schedule_offset: 剪枝开始的训练步数,0 表示立即开始。
示例代码:
import deepspeed
from deepspeed.compression.compress import init_compression, redundancy_clean
# 加载模型
model = ... # 假设已加载模型
model.to("cuda").eval()
# 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = init_compression(model, deepspeed_config)
# 执行稀疏剪枝
model = redundancy_clean(model, deepspeed_config)
2.4.2 行剪枝(Row Pruning)
说明:行剪枝将权重矩阵的某些行置零,适用于 Transformer 前馈网络(如 BERT 的 intermediate.dense 层),硬件加速效果较好。
JSON 配置:
{
"compression_training": {
"row_pruning": {
"enabled": true, // 是否启用行剪枝,布尔值,true 表示启用
"sparse_ratio": 0.4, // 行稀疏比率,浮点数,范围 [0, 1],表示移除的行比例
"modules": ["intermediate.dense"] // 应用行剪枝的模块名列表
}
}
}
JSON 配置解释:
enabled: 控制是否启用行剪枝。sparse_ratio: 移除行的比例,如 0.4 表示移除 40% 的行。modules: 指定应用行剪枝的模块名称,如intermediate.dense表示前馈网络的中间层。
示例代码:
import deepspeed
from deepspeed.compression.compress import init_compression, redundancy_clean
# 加载模型
model = ... # 假设已加载模型
model.to("cuda").eval()
# 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = init_compression(model, deepspeed_config)
# 执行行剪枝
model = redundancy_clean(model, deepspeed_config)
2.4.3 头剪枝(Head Pruning)
说明:头剪枝移除 Transformer 多头注意力机制中不重要的注意力头,适用于 NLP 模型。
JSON 配置:
{
"compression_training": {
"head_pruning": {
"enabled": true, // 是否启用头剪枝,布尔值,true 表示启用
"head_sparse_ratio": 0.3, // 注意力头稀疏比率,浮点数,范围 [0, 1],表示移除的头比例
"modules": ["attention.self"] // 应用头剪枝的模块名列表
}
}
}
JSON 配置解释:
enabled: 控制是否启用头剪枝。head_sparse_ratio: 移除注意力头的比例,如 0.3 表示移除 30% 的头。modules: 指定应用头剪枝的模块,如attention.self表示注意力层。
示例代码:
import deepspeed
from deepspeed.compression.compress import init_compression, redundancy_clean
# 加载模型
model = ... # 假设已加载模型
model.to("cuda").eval()
# 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = init_compression(model, deepspeed_config)
# 执行头剪枝
model = redundancy_clean(model, deepspeed_config)
2.4.4 通道剪枝(Channel Pruning)
说明:通道剪枝移除卷积网络或 Transformer 中的某些通道,适用于视觉模型或特定 Transformer 层。
JSON 配置:
{
"compression_training": {
"channel_pruning": {
"enabled": true, // 是否启用通道剪枝,布尔值,true 表示启用
"channel_sparse_ratio": 0.2, // 通道稀疏比率,浮点数,范围 [0, 1],表示移除的通道比例
"modules": ["conv1", "conv2"] // 应用通道剪枝的模块名列表
}
}
}
JSON 配置解释:
enabled: 控制是否启用通道剪枝。channel_sparse_ratio: 移除通道的比例,如 0.2 表示移除 20% 的通道。modules: 指定应用通道剪枝的模块名称,如conv1、conv2。
示例代码:
import deepspeed
from deepspeed.compression.compress import init_compression, redundancy_clean
# 加载模型
model = ... # 假设已加载模型
model.to("cuda").eval()
# 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = init_compression(model, deepspeed_config)
# 执行通道剪枝
model = redundancy_clean(model, deepspeed_config)
3. ZeroQuant 教程:高效且经济实惠的训练后量化
说明:ZeroQuant 是一种高效训练后量化(PTQ)方法,通过细粒度量化和层级知识蒸馏(LKD)实现 INT8 或 INT4/INT8 混合精度量化,保持高精度。
JSON 配置:
{
"compression_training": {
"zero_quant": {
"enabled": true, // 是否启用 ZeroQuant,布尔值,true 表示启用
"quantization_type": "symmetric", // 量化类型,字符串,可选 "symmetric" 或 "asymmetric"
"weight_bits": 8, // 权重量化位宽,整数,范围 [1, 8],如 8 表示 INT8
"activation_bits": 8, // 激活量化位宽,整数,范围 [1, 8],如 8 表示 INT8
"lkd_enabled": true, // 是否启用层级知识蒸馏,布尔值,true 表示启用
"lkd_layers": ["encoder.layer.0", "encoder.layer.1"] // 应用 LKD 的层名称列表
}
}
}
JSON 配置解释:
enabled: 控制是否启用 ZeroQuant。quantization_type: 量化类型,symmetric或asymmetric。weight_bits: 权重量化位宽,如 8 表示 INT8,4 表示 INT4。activation_bits: 激活量化位宽,通常为 8(INT8)。lkd_enabled: 是否启用层级知识蒸馏,增强量化精度。lkd_layers: 指定应用 LKD 的层名称,空列表表示所有层。
示例代码:
import deepspeed
from transformers import BertModel
import torch
# 加载预训练模型
model = BertModel.from_pretrained("bert-base-uncased").to("cuda")
model.eval()
# DeepSpeed 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = deepspeed.compression.init_compression(model, deepspeed_config)
# 执行 ZeroQuant 量化
model = deepspeed.compression.redundancy_clean(model, deepspeed_config)
# 推理
data = torch.rand(4, 128, device="cuda")
outputs = model(data)
4. XTC 教程:用于极端压缩的简单而有效的压缩管道
说明:XTC(eXTreme Compression)结合层减少和超低位量化(如 1 位权重),可将模型大小压缩 32x 甚至 50x,同时保持高精度。
JSON 配置:
{
"compression_training": {
"xtc": {
"enabled": true, // 是否启用 XTC,布尔值,true 表示启用
"layer_reduction": {
"keep_number_of_layers": 5 // 保留的层数,整数,需小于原始层数
},
"weight_quantization": {
"enabled": true, // 是否启用权重量化,布尔值,true 表示启用
"quantization_type": "binary", // 量化类型,字符串,"binary" 表示 1 位量化
"target_bits": 1 // 目标量化位宽,整数,1 表示二值化
},
"activation_quantization": {
"enabled": true, // 是否启用激活量化,布尔值,true 表示启用
"target_bits": 8 // 目标量化位宽,整数,8 表示 INT8
}
}
}
}
JSON 配置解释:
enabled: 控制是否启用 XTC。layer_reduction:keep_number_of_layers: 保留的层数,如 5。
weight_quantization:enabled: 控制权重量化。quantization_type:"binary"表示 1 位二值化。target_bits: 目标位宽,1 表示二值化。
activation_quantization:enabled: 控制激活量化。target_bits: 目标位宽,8 表示 INT8。
示例代码:
import deepspeed
from transformers import BertModel
import torch
# 加载预训练模型
model = BertModel.from_pretrained("bert-base-uncased").to("cuda")
model.eval()
# DeepSpeed 压缩初始化
deepspeed_config = "deepspeed_config.json"
model = deepspeed.compression.init_compression(model, deepspeed_config)
# 执行 XTC 压缩
model = deepspeed.compression.redundancy_clean(model, deepspeed_config)
# 推理
data = torch.rand(4, 128, device="cuda")
outputs = model(data)
运行示例
运行 DeepSpeed 压缩的示例脚本:
# 安装依赖
cd DeepSpeedExamples/model_compression/bert
pip install -r requirements.txt
# 运行 XTC 压缩
bash bash_script/XTC/run_glue_no_trainer.py \
--model_name_or_path bert-base-uncased \
--deepspeed_config deepspeed_config.json
5. 代码示例:压缩与推理
本示例使用 Hugging Face 的 IMDb 数据集(情感分类)和 DistilBERT 模型。也可以替换为其他模型(如 LLaMA)或数据集。
以下是一个完整的示例,展示如何使用 DeepSpeed Compression 压缩 DistilBERT 模型(层减少 + INT8 量化),并结合 deepspeed.init_inference() 进行高效推理。
import deepspeed
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from deepspeed.compression import init_compression, redundancy_clean
from datasets import load_dataset
import os
# Environment setup
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29500"
# Load model and tokenizer
model_name = "distilbert-base-uncased"
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# DeepSpeed compression config
ds_config = {
"compression_training": {
"layer_reduction": {
"enabled": True,
"keep_number_layer": 6, #减少到 6 层 (目标层数)
"module_name_prefix": "bert.encoder.layer",
"teacher_layer": [0, 2, 4, 6, 8, 10], # 教师模型层选择(从教师模型中选择的层索引)
"other_module_name": ["bert.embeddings", "bert.pooler", "classifier"] # 保留的非 Transformer 层(如嵌入层、分类层)
},
"weight_quantization": {
"enabled": True,
"quantize_type": "symmetric", # 对称量化(symmetric)或非对称量化(asymmetric)
"num_bits": 8 # 量化位数(如 8 表示 INT8,4 表示 INT4)
}
},
"fp16": {"enabled": True},
"zero_optimization": {
"stage": 3,
"offload_param": {"device": "cpu", "pin_memory": True}
}
}
# Apply compression
model = init_compression(model, ds_config)
# Load dataset for fine-tuning (optional)
dataset = load_dataset("imdb")
def preprocess(examples):
return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
tokenized_dataset = dataset.map(preprocess, batched=True)
# Fine-tune compressed model (simplified)
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir="./compressed_output",
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
num_train_epochs=1,
deepspeed=ds_config
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset["train"]
)
trainer.train()
# Clean compression
model = redundancy_clean(model)
# Initialize DeepSpeed inference
ds_engine = deepspeed.init_inference(
model=model,
mp_size=2, // 2 GPUs
dtype=torch.float16,
replace_with_kernel_inject=True,
enable_cuda_graph=True
)
# Inference
text = "This movie is fantastic!"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to("cuda")
with torch.no_grad():
outputs = ds_engine(inputs)
logits = outputs.logits
pred = torch.argmax(logits, dim=-1)
print(f"Text: {text}, Prediction: {'Positive' if pred.item() == 1 else 'Negative'}")
6. 运行步骤
5.1 保存代码和配置
- 保存 DeepSpeed 配置文件为
ds_config.json。 - 保存上述代码为
compress_infer.py。
5.2 运行命令
使用 deepspeed 命令运行(假设有 2 块 GPU):
deepspeed --num_gpus=2 compress_infer.py
5.3 输出
- 训练日志:显示压缩和微调进度。
- 推理结果:
Text: This movie is fantastic! Prediction: Positive
7. 关键实现细节
7.1 压缩流程
- 初始化压缩:
init_compression根据ds_config修改模型结构(如减少层数、量化权重)。 - 微调:压缩后的模型通过任务特定数据(如 IMDb)微调,恢复精度。
- 清理压缩:
redundancy_clean优化模型结构,确保推理兼容性。 - 推理:
deepspeed.init_inference加载压缩模型,应用张量并行和优化内核。
7.2 与 deepspeed.init_inference() 的结合
- 压缩模型支持:压缩后的模型直接兼容
deepspeed.init_inference,支持 INT8 权重和激活。 - 推理优化:
- 张量并行(
mp_size)分区压缩模型权重。 - FP16/BF16(
dtype)进一步降低内存。 - 优化内核(
replace_with_kernel_inject)加速 INT8 计算。 - CUDA Graph(
enable_cuda_graph)减少固定形状输入的延迟。
- 张量并行(
7.3 ZeroQuant 特点
- 零成本量化:无需量化感知训练(QAT),直接后训练量化(PTQ)。
- 层级知识蒸馏(LKD):通过逐层蒸馏保留量化模型精度。
- 支持 INT4/INT8:如 BERT 和 GPT-3 模型的混合精度量化。
7.4 XTC 特点
- 极致压缩:结合层减少、二值化和知识蒸馏,压缩比高达 32x。
- 模块化 API:通过压缩组合器简化多方法协同应用。
8.学习资源
- DeepSpeed 文档:https://www.deepspeed.ai/docs/
- 压缩教程:https://www.deepspeed.ai/tutorials/compression/
- 中文压缩教程:https://deepspeed.org.cn/tutorials/model-compression/#1-general-tutorial
- ZeroQuant 论文:https://arxiv.org/abs/2206.01861
- GitHub 源码:https://github.com/microsoft/DeepSpeed
- DeepSpeed 示例:https://github.com/microsoft/DeepSpeedExamples
更多推荐


所有评论(0)