使用 DeepSpeed 进行分布式训练

上一页
下一页

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

基础分布式训练

训练机器学习模型包含许多部分，通常会将这些部分分布到多个计算资源（例如GPU或集群节点）上。扩展数据和计算通常面临一项重要挑战，导致大型模型的轻松高效训练。

为什么选择分布式训练？

在处理大型深度学习模型时，需要考虑分布式训练的关键原因如下：

  • 可扩展性 - 在单个GPU上训练具有数千万甚至数十亿参数的超大型模型非常困难。通过使用分布式训练，可以将此过程扩展到多个GPU上。
  • 更快的收敛速度 - 将训练过程分散到多个GPU上可以加快收敛过程，从而加快模型开发速度。
  • 资源效率 - 此类训练将充分利用您的可用硬件，从而节省时间和金钱。
  • 数据并行 - 将一个模型分布在多个GPU上，每个GPU处理数据集的不同批次。
  • 模型并行 - 模型在多个GPU上并行化；每个GPU计算模型操作的一部分。
  • 混合并行 - 混合数据和模型并行。换句话说，将数据分割到GPU上，然后进一步分割模型。

数据并行

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

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

使用 DeepSpeed 的数据并行示例

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

import torch
import deepspeed

# Define a simple neural network model
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)

# Initialize DeepSpeed configuration
deepspeed_config = {
    "train_batch_size": 64,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    }
}

# Initialize model
model = SimpleModel()

# Initialize DeepSpeed for distributed data parallelity
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model
)

# Dummy data
inputs = torch.randn(64, 784)
labels = torch.randint(0, 10, (64,))

# Forward pass
outputs = model_engine(inputs)
loss = torch.nn.functional.cross_entropy(outputs, labels)

# Backward pass and optimization
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

# Define a simple pipeline model
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))

# Two GPUs and two layers in a pipeline paradigm.
layers = [
    LayerSpec(SimpleLayer, 784, 128),
    LayerSpec(SimpleLayer, 128, 10)
]
# We create a pipeline model, specifying the number of stages - 2
pipeline_model = PipelineModule(layers=layers, num_stages=2)

# Initialize DeepSpeed for model parallelism
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=pipeline_model
)

# Dummy inputs
inputs = torch.randn(64, 784)

# Forward pass through pipeline
outputs = model_engine(inputs)

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

零冗余优化器 (ZeRO)

DeepSpeed 最显著的特性也许是零冗余优化器 (Zero Redundancy Optimizer)，更方便地称为 ZeRO，它旨在解决模型训练的内存消耗问题。它将各种状态分布在不同的GPU上，从而更有效地利用内存：优化器、梯度和参数。

ZeRO 包括三个阶段：

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

零冗余优化器示例

以下是 Python 中零冗余优化器的简单示例：

import torch
import deepspeed

# Use ZeRO optimization to define the model and DeepSpeed settings
deepspeed_config = {
    "train_batch_size": 64,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    },
    "zero_optimization": {
        "stage": 2 # Toggle gradient partitioning using ZeRO Stage 2
    }
}

# Initialize model
model = SimpleModel()

# Initialize DeepSpeed with ZeRO optimization
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model
)

# Forward pass
inputs = torch.randn(64, 784)
outputs = model_engine(inputs)

# Backward pass and optimization
model_engine.backward(outputs)
model_engine.step()

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

跨多个GPU和节点扩展模型

DeepSpeed 通过利用混合并行策略和 DeepSpeed 的高级通信层来实现最佳扩展，从而跨多个 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
# Training on multiple GPUs
if torch.distributed.get_rank() == 0:
    print("Training on multiple GPUs with DeepSpeed")
# Initialize DeepSpeed with ZeRO optimization for multi-GPU
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=deepspeed_config
)
# Training loop
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 的特性将同时确保分布式训练在增长时保持可访问性和性能增强。

打印页面
上一页
下一页
广告