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 info: version=0.6.0, git-hash=unknown, git-branch=unknown
[INFO] Initializing model parallel group with size 1
[INFO] Initializing optimizer with 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都会给你类似这样的结果

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的高级功能，从而推动你的深度学习项目。

打印页面