使用 DeepSpeed 进行模型训练

深度学习模型变得越来越大、越来越复杂，使得训练过程更加难以有效进行。这就是微软的 DeepSpeed-深度学习优化库发挥作用的地方。该库旨在训练大型模型；它还拥有一系列旨在优化内存、提高计算效率和整体训练性能的功能。本章结束时的目标将包括使用 DeepSpeed 进行训练、查看设置优化功能的配置文件，并提供使用此强大工具训练流行模型的一些示例。

使用 DeepSpeed 进行深度学习模型训练

训练深度学习模型是一项计算密集型任务，尤其是在处理大型数据集和复杂架构时。 DeepSpeed 专为应对这一挑战而构建，它在一个框架中提供了一组功能，包括混合精度训练、ZeRO(零冗余优化器)和梯度累积，可确保极高的效率，同时扩大模型训练规模，而不必成倍地扩大计算资源。

现在我们将开始将 DeepSpeed 实现到一个简单的模型训练管道中。

步骤 1:模型和数据集

假设一个简单的 PyTorch 模型正在解决回归问题:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 一个简单的回归模型
class RegressionModel(nn.Module):
    def __init__(self):
        super(RegressionModel, self).__init__()
        self.fc1 = nn.Linear(10, 50)
self.fc2 = nn.Linear(50, 1)
def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 生成合成数据
inputs = torch.randn(1000, 10)
targets = torch.randn(1000, 1)
dataset = TensorDataset(inputs, target)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

model = RegressionModel()

第 2 步:添加 DeepSpeed

下一步是将 DeepSpeed 添加到您的配置文件中以启用训练优化。

DeepSpeed 配置文件

DeepSpeed 配置文件是 JSON 文件，它指定了优化模型训练的多个参数。示例如下:

{
    "train_batch_size": 32,
    "fp16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 1,
        "allgather_partitions": true,
        "reduce_scatter": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true
    },
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 3e-7
        }
    }
}

将上述文本保存到项目文件夹中名为 ds_config.json 的文件中。

步骤 3:DeepSpeed 初始化

事情从这里开始变得有趣。设置好配置文件后，您就可以在训练脚本中初始化 DeepSpeed，如下所示:

import deepspeed

# 初始化 DeepSpeed
ds_config_path = "ds_config.json"
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=ds_config_path
)

输出

运行上述代码将使用下面指定的配置初始化 DeepSpeed −

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

使用 DeepSpeed 的功能优化训练

DeepSpeed 带有一组可以优化模型训练的功能。我们将在这里讨论一些关键功能。

混合精度训练 −它以 16 位浮点表示形式训练模型，因此需要更少的内存，从而实现更快的计算。
ZeRO 优化 − 零冗余优化器 (ZeRO) 可以通过在数千个 GPU 上划分模型状态来大幅减少大型模型的内存占用。您可以使用 zero_optimization 部分中的 stage 参数值来控制优化程度。
梯度累积 − 此功能允许增加有效批次大小，而无需按比例增加 GPU 内存。您可以通过在配置文件中设置 gradient_accumulation_steps 的值来启用梯度累积。
激活检查点 − 这种方法是一种计算与内存节省方法，因为它以重新计算后向传递中的一些激活为代价来节省内存。这意味着它可以减少训练时的整体内存消耗。

这些功能可以以各种方式组合，具体取决于最适合您的特定要求的方式。

使用 DeepSpeed 训练 BERT 模型的示例

为了展示 DeepSpeed 的强大功能，以 BERT − 等著名模型的训练为例来自 Transformers 的双向编码器表示。

步骤 1:准备和加载 BERT 模型

您可以使用 Hugging Face Transformers 库轻松加载预先训练的 BERT 模型 −

from transformers import BertForSequenceClassification, BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained("bert-base-uncased")

# 示例数据
inputs = tokenizer("DeepSpeed makes BERT training efficient!", return_tensors="pt")
labels = torch.tensor([1]).unsqueeze(0)

# Dataloader
dataloader = DataLoader([(inputs, labels)], batch_size=1)

第 2 步:添加 DeepSpeed 集成

与之前一样，我们通过初始化您的模型和配置文件来添加 DeepSpeed 集成 −

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

步骤 3:运行模型

模型如下 −

for batch in dataloader:
        inputs, labels = batch
        outputs = model_engine(**inputs)
loss = nn.CrossEntropyLoss()(outputs.logits,labels)

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

输出

使用 DeepSpeed 训练 BERT 将输出每个时期的损失，确保模型得到有效训练 −

Epoch 1, Loss: 0.6785
Epoch 2, Loss: 0.5432
Epoch 3, Loss: 0.4218

使用 DeepSpeed 处理大型数据集

大型数据集带来的问题远远超出了模型架构的范围。如何在处理大量数据的同时有效地管理内存和计算资源将使您免于瓶颈。DeepSpeed 通过其在数据处理领域的高级功能解决了这些挑战。

1. 动态数据加载

DeepSpeed 执行数据的动态加载，从而只将训练期间一次使用的批次加载到内存中。这减少了内存占用，因此允许您在更大的数据集上进行训练，而不一定需要更强大的硬件。除此之外，您将保持内存使用量最小；因此，您可以最大限度地减少数据输入/输出操作所花费的时间，从而提高整体训练速度。

2. 数据并行性

DeepSpeed 支持的另一项重要功能是数据并行性。它支持跨许多 GPU 的本地分布式数据。因此，可以一次处理不同的批次。这种并行将加快训练过程。它可以有效地占用 GPU 资源。因此，在实践中，使用 DeepSpeed 将数据并行性应用于训练管道并不困难，因为它已集成到 PyTorch 的 DataLoader 中。

3. 内存高效的数据改组

大型数据集通常需要改组以避免过度拟合和根据数据的排序方式进行模式学习。但是，这对于大型数据集来说非常耗费内存。DeepSpeed 使用非常节省内存的算法优化了此过程，能够在不增加大量内存的情况下提供有效的改组。这确保了在大型数据集上，训练将顺利而高效。

4. 数据增强支持

数据增强通常包括某些通过修改现有数据来人为增加数据集大小的方法。 DeepSpeed 支持即时数据增强，这意味着不必将增强数据存储在内存中，而是可以在训练期间即时执行数据增强。这可以进一步减少内存压力，并提供更广泛的数据增强技术利用。

5. 批量大小扩展

借助 DeepSpeed 梯度累积和 ZeRO 优化，即使在处理大量数据集时也可以扩大批量大小。更大的批量大小有时可以提高模型收敛性和训练稳定性。启用 DeepSpeed 后，可以在管理 GPU 内存需求的情况下扩展批量大小；因此，您的模型应该能够有效地在大数据集上进行训练。

上述 DeepSpeed 功能在这方面有所帮助，因为它能够管理大型数据集，从而使您可以设计和训练不受硬件限制的高性能模型。无论您是在非常大的文本语料库上训练模型，还是处理超高分辨率的图像，DeepSpeed 处理数据的这一功能都能让您的训练流程保持优化和可扩展。

总结

DeepSpeed 允许为深度学习模型提供有效的训练框架，尤其是在扩展规模和复杂性方面。因此，学习如何使用混合精度训练、ZeRO 优化和激活检查点的高级功能是增加价值优化流程的方式。本章包含有关使用 DeepSpeed 进行模型训练的信息，为 DeepSpeed 准备环境、DeepSpeed 的配置以及运行训练过程。有了这个工具和技术，您现在可以以更好的性能和更低的资源消耗处理大规模深度学习项目。

DeepSpeed 教程

DeepSpeed 有用资源