使用 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

# A simple regression model
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

# Generating synthetic data
inputs = torch.randn(1000, 10)
targets = torch.randn(1000, 1)
dataset = TensorDataset(inputs, targets)
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

# Initialize 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 info: version=0.6.0, git-hash=unknown, git-branch=unknown
[INFO] Initializing model parallel group with size 1
[INFO] Initialize optimizer with DeepSpeed Zero Optimizer

使用 DeepSpeed 的功能优化训练

DeepSpeed 带有一系列可以优化模型训练的功能。我们将在本文中讨论其中一些关键功能。

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

可以根据您的特定需求以各种方式组合这些功能。

使用 DeepSpeed 训练 BERT 模型的示例

为了展示 DeepSpeed 的强大功能,我们以训练著名的 BERT 模型(来自 Transformer 的双向编码器表示)为例。

步骤 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")

# Sample data
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 的损失,确保模型正在高效地训练 -

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 配置以及运行训练过程的信息。有了此工具和技术,您现在可以使用更好的性能和更低的资源消耗来处理大型深度学习项目。

广告