DeepSpeed
LLM（十二）：DeepSpeed Inference 在 LLM 推理上的优化探究
Zero Redundancy Optimizer (ZeRO)是通过划分优化器状态、梯度、模型参数到不同GPU上达到降低大模型训练显存需求，提升显存利用率目的。
主要分为3个Stage：
- ZeRO Stage1:对优化器状态分区（Adam优化器的话，包含1阶和2阶动量参数、32bit的模型参数备份）
- ZeRO Stage2：对优化器状态和梯度进行分区（32bit的梯度）
- ZeRO Stage3:对优化器状态、梯度以及模型参数(16bit)都进行分区
另外，ZeRO Stage 3还支持ZeRO-Infinity优化手段，将参数offload到CPU内存和硬盘上，进一步减小显存占用
DeepSpeed和Accelerate
DeepSpeed 是由 Microsoft 开发的一个深度学习优化库，旨在提高大规模模型训练的效率。它提供了多种优化技术，包括混合精度训练、分布式训练、数据并行、模型并行和高效的梯度累积等。 主要特点：
- 分布式训练: 支持数据并行、模型并行（包括管道并行和张量并行），使得训练超大规模模型成为可能。
- 优化技术: 提供优化算法，如 ZeRO（Zero Redundancy Optimizer）以减少显存占用、加速训练速度。
- 混合精度训练: 自动支持 FP16 和 BFLOAT16 精度训练，以减少计算开销和内存占用。
- 高效的梯度累积: 提供高效的梯度累积方法来处理超大批量训练。
- 弹性训练: 支持弹性训练，允许动态添加或移除计算资源。 功能非常的强大、丰富，但是配置起来比较复杂，需要一定的深度学习知识。
Accelerate 是由 Hugging Face 开发的一个库，旨在简化分布式训练的设置。它提供了一种简洁的方式来配置和管理多 GPU 和 TPU 环境，支持数据并行和模型并行。 主要特点
- 简化配置: 提供统一的接口来处理多 GPU 和 TPU 环境的配置，简化了分布式训练的复杂性。
- 支持多种后端: 可以与不同的深度学习后端（如 PyTorch 和 TensorFlow）集成。
- 集成 DeepSpeed: 可以与 DeepSpeed 集成，利用 DeepSpeed 的高级功能进行训练。
- 简化分布式训练: 自动处理分布式训练的设置和同步问题，使用户能够专注于模型和数据
功能丰富，而且使用非常简单，但是配置不是非常精细。 联系
- 集成能力: Accelerate 可以与 DeepSpeed 集成，利用 DeepSpeed 的高级功能来优化训练。通过 Accelerate 的 DeepSpeedPlugin，可以在 Accelerate 的框架下使用 DeepSpeed 进行训练。
- 共同目标: 两者都旨在提高大规模模型训练的效率和简化配置，但 DeepSpeed 提供了更多的优化功能，而 Accelerate 注重于简化配置和多后端支持。
代码示例：
{
"compute_environment": "LOCAL_MACHINE",
"debug": true,
"distributed_type": "MULTI_GPU",
"downcast_bf16": false,
"enable_cpu_affinity": false,
"machine_rank": 0,
"main_training_function": "main",
"mixed_precision": "no",
"num_machines": 1, ## 几个机器
"num_processes": 2, ## 几个GPU
"rdzv_backend": "static",
"same_network": true,
"tpu_use_cluster": false,
"tpu_use_sudo": false,
"use_cpu": false
} 可以运行的代码例子
from accelerate import Accelerator, DeepSpeedPlugin
import torch
from torch.utils.data import DataLoader, TensorDataset
class SimpleNet(torch.nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SimpleNet, self).__init__()
self.fc1 = torch.nn.Linear(input_dim, hidden_dim) #.to("cuda:0")
self.fc2 = torch.nn.Linear(hidden_dim, output_dim) #.to("cuda:1")
def forward(self, x):
# x.to("cuda:0")
x = torch.relu(self.fc1(x))
# x.to("cuda:1")
x = self.fc2(x)
return x
if __name__ == "__main__":
input_dim = 10
hidden_dim = 20
output_dim = 2
batch_size = 64
data_size = 10000
input_data = torch.randn(data_size, input_dim)
labels = torch.randn(data_size, output_dim)
dataset = TensorDataset(input_data, labels)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
model = SimpleNet(input_dim, hidden_dim, output_dim)
deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_clipping=1.0)
accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
optimization = torch.optim.Adam(model.parameters(), lr=0.00015)
crition = torch.nn.MSELoss()
model, dataloader, optimization = accelerator.prepare(model, dataloader, optimization)
for epoch in range(1000):
model.train()
for batch in dataloader:
inputs, labels = batch
outputs = model(inputs)
loss = crition(outputs, labels)
optimization.zero_grad()
accelerator.backward(loss)
optimization.step()
print(f"Epoch {epoch} loss: {loss.item()}")
accelerator.save(model.state_dict(), "model.pth")