使用 DeepSpeed 进行分布式训练

随着模型大小和数据集大小的增加，单 GPU 训练在大多数情况下变得效率低下或不可能。因此，基于分布式训练，模型可以轻松地从单个 GPU 扩展到多个 GPU 和节点。结合优化这种训练方法，微软的 DeepSpeed 是装备最精良的框架之一。它能够处理大型模型，并通过数据并行、模型并行和零冗余优化器 (ZeRO) 等突出技术减少内存开销。

基本分布式训练

训练机器学习模型有很多部分，通常，它会分成多个部分，分布在许多计算资源上，例如 GPU 或集群节点。扩大数据和计算通常面临重要挑战，导致大型模型的训练变得简单而高效。

为什么要进行分布式训练?

以下是在处理大型深度学习模型时考虑进行分布式训练的主要原因 −

可扩展性 − 在单个 GPU 上训练具有数千万甚至数十亿个参数的大型模型极其困难。可以使用分布式训练跨多个 GPU 扩展此过程。
更快的收敛 − 在多个 GPU 之间分散训练过程可加速收敛过程，从而加快模型开发速度。
资源效率 −这种训练可以最大限度地利用可用硬件，从而节省时间和金钱。
数据并行 − 当一个模型分布在多个 GPU 上，并且每个 GPU 处理不同批次的数据集时，就会出现这种情况。
模型并行 − 该模型在多个 GPU 上并行化；每个 GPU 计算模型操作的部分内容。
混合并行 − 混合数据和模型并行。换句话说，在 GPU 上拆分数据并进一步拆分模型。

数据并行

DeepSpeed 通过提供适应性模型和数据并发来促进分布式训练。让我们深入探讨这些。

使用数据并行时，每个 GPU 或工作器都会接收一部分数据进行处理。然后，它会在处理后对这些结果进行平均，以更新模型权重。因此，可以使用更大的批量大小进行训练，而不会耗尽内存。

使用 DeepSpeed 进行数据并行的示例

以下是一个简单的 Python 示例，用于展示使用 DeepSpeed 进行数据并行 −

import torch
import deepspeed

# 定义一个简单的神经网络模型
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = torch.nn.Linear(784, 128)
        self.fc2 = torch.nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# 初始化 DeepSpeed 配置
deepspeed_config = {
    "train_batch_size": 64,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    }
}

# 初始化模型
model = SimpleModel()

# 初始化 DeepSpeed 以实现分布式数据并行
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model
)

# 虚拟数据
inputs = torch.randn(64, 784)
labels = torch.randint(0, 10, (64,))

# 前向传递
outputs = model_engine(inputs)
loss = torch.nn. functional.cross_entropy(outputs, labels)

# 后向传递和优化
model_engine.backward(loss)
model_engine.step()

神经网络现在将在多个 GPU 上进行训练；每个 GPU 负责一部分数据。

模型并行性

模型并行性涉及将模型拆分到多个 GPU 上。当单个模型无法装入单个 GPU 的内存时，这会变得很有帮助。

使用 DeepSpeed 实现模型并行

它将模型拆分到多个 GPU 上，其中模型的不同部分可以同时在不同的 GPU 上执行。

使用 DeepSpeed 实现模型并行的示例

以下是一个简单的 Python 程序，用于展示使用 DeepSpeed 实现模型并行的工作原理 −

import torch
import deepspeed
from deepspeed.pipe import PipelineModule, LayerSpec

# 定义一个简单的管道模型
class SimpleLayer(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super(SimpleLayer, self).__init__()
        self.fc = torch.nn.Linear(input_size, output_size)

    def forward(self, x):
        return torch.relu(self.fc(x))

# 管道范例中的两个 GPU 和两层。
layers = [
    LayerSpec(SimpleLayer, 784, 128),
    LayerSpec(SimpleLayer, 128, 10)
]
# 我们创建一个管道模型，指定阶段数 - 2
pipeline_model = PipelineModule(layers=layers, num_stages=2)

# 初始化 DeepSpeed 以实现模型并行
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=pipeline_model
)

# 虚拟输入
inputs = torch.randn(64, 784)

# 前向传递管道
outputs = model_engine(inputs)

这将分阶段处理跨多个 GPU 的前向传递。第一个 GPU 将处理到第一层，而第二个 GPU 将处理到倒数第二层。

零冗余优化器 (ZeRO)

DeepSpeed 最突出的功能可能是零冗余优化器，更方便地称为 ZeRO，旨在解决模型训练的内存消耗问题。它将各种状态拆分到不同的 GPU，从而更有效地使用内存:优化器、梯度和参数。

ZeRO 包括三个阶段 −

第 1 阶段 −对优化器的状态进行分区。
第 2 阶段 − 对梯度状态进行分区。
第 3 阶段 − 对参数状态进行分区。

零冗余优化器示例

以下是 python 中零冗余优化器的简单示例 −

import torch
import deepspeed

# 使用 ZeRO 优化定义模型和 DeepSpeed 设置
deepspeed_config = {
    "train_batch_size": 64,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    },
    "zero_optimization": {
        "stage": 2 # Toggle gradient partitioning using ZeRO Stage 2
    }
}

# 初始化模型
model = SimpleModel()

# 使用 ZeRO 优化初始化 DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model
)

# 前向传递
inputs = torch.randn(64, 784)
outputs = model_engine(inputs)

# 后向传递和优化
model_engine.backward(outputs)
model_engine.step()

此代码在 ZeRO Stage 2 上运行，该阶段是跨 GPU 划分的梯度状态，可减少训练期间的内存消耗。

跨多个 GPU 和节点扩展模型

DeepSpeed 通过利用并行策略与 DeepSpeed 的高级通信层混合来实现跨多个 GPU 和节点扩展模型，以实现最佳扩展。

使用多个 GPU 和节点的扩展示例节点

NCCL 后端用于 GPU 间通信，并将训练扩展到多个 GPU 和节点。我们可以进行以下调用以使用在多个 GPU 和节点上运行的 DeepSpeed:

要使用 DeepSpeed 在多个 GPU 和节点上运行，您可以使用以下命令:

deepspeed --num_nodes 2 --num_gpus 8 train.py

这总共使用 8 个 GPU 和 2 个节点进行训练。

使用 DeepSpeed 在多个 GPU 上进行训练的示例

以下示例演示了如何使用 DeepSpeed 在多个 GPU 上进行训练 −

import deepspeed
# 在多个 GPU 上进行训练
if torch.distributed.get_rank() == 0:
    print("Training on multiple GPUs with DeepSpeed")
# 使用 ZeRO 优化初始化 DeepSpeed 以适应多 GPU
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=deepspeed_config
)
# 训练循环
for batch in train_loader:
    inputs, labels = batch
    outputs = model_engine(inputs)
    loss = torch.nn.functional.cross_entropy(outputs, labels)
    model_engine.backward(loss)
    model_engine.step()

此代码使用 DeepSpeed 在各种 GPU 上高效地训练模型，并采用 ZeRO 等方法进行优化。

总结

DeepSpeed 已得到强大的开发，可用于扩展和优化深度学习模型中的分布式训练。通过集成 ZeRO 以在多个 GPU 和节点上进一步扩展，结合数据并行和模型并行，DeepSpeed 可以完全解决高效训练大型模型的所有挑战。这意味着，DeepSpeed 的功能将同时确保分布式训练在增长过程中保持可访问性和性能增强。

DeepSpeed 教程

DeepSpeed 有用资源