开始使用 DeepSpeed

随着训练计算成本的上升,深度学习模型变得越来越复杂。DeepSpeed 由微软开发,能够以最少的资源高效地训练大规模模型。本章将带您了解使用 DeepSpeed 启动和运行模型的基本步骤,从安装和设置环境到运行第一个模型。

安装 DeepSpeed

在进一步深入了解 DeepSpeed 的细节之前,我们需要做的第一件事是安装库。使用 pip,这很容易实现 −

pip install deepspeed

安装时,您可能会看到类似下面的结果 −

Collecting deepspeed
Downloading deepspeed-0.6.0-py3-none-any.whl (696 kB)
|████████████████████████████████| 696 kB 3.2 MB/s 
Collecting torch
Downloading torch-1.9.1-cp38-cp38-manylinux1_x86_64.whl (804.1 MB)
|████████████████████████████████|
deepspeed-0.6.0 torch-1.9.1 installed successfully

如果您愿意,您还可以克隆 GitHub 存储库并从源代码安装 −

git clone https://github.com/microsoft/DeepSpeed.git
cd DeepSpeed
pip install .

这将为您提供最新功能,这些功能可能尚未在稳定版本中发布。

环境设置

安装 DeepSpeed 后,必须设置环境。首先,确保所有必需的依赖项都存在。

创建一个虚拟环境来管理依赖项 −

python -m venv deepspeed-env
source deepspeed-env/bin/activate  # On Windows, use 'deepspeed-env\Scripts\activate'

如果尚未安装 PyTorch,请先安装 −

pip install torch torchvision torchaudio

此外,根据您的使用情况,您可能需要 CUDA 或其他类型的 GPU 加速。如果您使用的是配备 GPU 的机器,则安装 CUDA 版本的 PyTorch 非常简单,只需在终端中运行以下命令即可 −

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

这将确保设置正确,以便 DeepSpeed 能够使用您机器的所有硬件功能。

基本概念和术语

在运行您的第一个模型之前,让我们先介绍一下您在 DeepSpeed 中经常遇到的一些基本概念和术语。

  • 优化器 − DeepSpeed 目前支持多个优化器,可用于优化大型模型的训练。优化器在训练模型时处理梯度更新。
  • 调度程序− 调度程序在训练期间更新学习率。默认情况下,DeepSpeed 集成了所有 PyTorch 调度程序,并进一步提供了为大型模型开发的额外自定义调度程序。
  • 零冗余优化器 (ZeRO) − 这是一种内存优化技术,通过将模型状态划分到多个 GPU 上来减少大型模型的内存占用。
  • 累积梯度 − 这可以通过在模型权重更新之前对多次迭代中的梯度求和来促进使用比允许 GPU 内存更大的批量大小。
  • 检查点激活 −这样可以节省一些内存,但需要额外的计算,在反向传播期间重新计算前向传递激活。

理解这些概念应该会为您提供足够的背景信息,以便您了解 DeepSpeed 中的大多数高级功能并自定义您的训练流程。

使用 DeepSpeed 运行您的第一个模型

现在您的环境已经设置好,并且您已经熟悉了基本术语,让我们运行一个简单的 DeepSpeed 模型。我们将首先创建一个基本的 PyTorch 模型,然后将 DeepSpeed 添加到其中以查看性能提升。

步骤 1:创建一个简单的 PyTorch 模型

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

class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(10, 50) # input layer (10) -> hidden layer (50)
self.fc2 = nn.Linear(50, 1) # hidden layer (50) -> output layer (1)

def forward(self, x):
x = torch.relu(self.fc1(x)) # hidden layer activation function
x = self.fc2(x)
return x

model = SimpleModel()

第 2 步:实现 DeepSpeed

现在,让我们重构代码,使其与 DeepSpeed 兼容。我们将使用 DeepSpeed 和一些基本配置初始化模型。

import deepspeed

ds_config = {
   "train_batch_size": 32,
   "fp16": {
      "enabled": True
   },
   "zero_optimization": {
      "stage": 1
   }
}

model_engine, optimizer, _, _ = deepspeed.initialize(
   model=model,
   model_parameters=model.parameters(),
   config=ds_config
)

输出

如果一切顺利,DeepSpeed 将初始化并打印出配置设置 −

[INFO] DeepSpeed 信息:version=0.6.0,git-hash=unknown,git-branch=unknown
[INFO] 初始化大小为 1 的模型并行组
[INFO] 使用 DeepSpeed Zero Optimizer 初始化优化器

步骤 3:训练模型

此时,您现在应该能够使用 DeepSpeed 训练您的模型。下面是一个示例训练循环。

for epoch in range(5) − 
inputs = torch.randn(32, 10)
labels = torch.randn(32, 1)

model_engine.train()
outputs = model_engine(inputs)
loss = nn.MSELoss()(outputs, labels)

model_engine.backward(loss)
model_engine.step()
print(f'Epoch {epoch+1}, Loss: {loss.item()}')

输出

每个时期都会给你一个类似这样的结果:

Epoch 1, Loss: 0.4857
Epoch 2, Loss: 0.3598
Epoch 3, Loss: 0.2893
Epoch 4, Loss: 0.2194
Epoch 5, Loss: 0.1745

第 4 步:保存模型

最后,您可以保存迄今为止训练过的模型 −

model_engine.save_checkpoint('./checkpoint', epoch=5)

输出

[INFO] Saving model checkpoint to ./checkpoint

DeepSpeed 的高级功能

现在让我们来看看 DeepSpeed 的一些高级功能,对 DeepSpeed 有一个基本的了解。这些高级功能的实现是为了处理训练大型模型的复杂性、减少内存消耗和提高计算效率。

  • 混合精度训练 FP16 − DeepSpeed 中模型训练速度快的原因之一是它支持使用半精度进行混合精度训练。
  • ZeRO 优化阶段 − DeepSpeed 拥有一项称为 ZeRO 的改变游戏规则的技术,它通过在多个 GPU 上划分模型状态来减少内存。
  • 梯度累积 − DeepSpeed 支持的另一种策略是梯度累积,它可以模拟更大的批量大小而不需要更多的 GPU 内存。
  • 卸载 −即使对于非常大的模型,ZeRO Stage 3 提供的优化也可能不够。

总结

开始使用 DeepSpeed 的主要步骤是安装库、设置环境、了解一些基本概念以及运行您的第一个模型。DeepSpeed 允许以更高的效率在更高的内存和更短的总体训练时间内训练大型模型。本基础章节将使您能够进一步了解 DeepSpeed 的高级功能,以推动您的深度学习项目。