使用 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 的特性將同時確保分散式訓練在增長時保持可訪問性和效能增強。

列印頁面
上一頁
下一頁
廣告
© . All rights reserved.