使用 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 配置以及執行訓練過程的資訊。有了此工具和技術,您現在可以使用更好的效能和更低的資源消耗來處理大型深度學習專案。

廣告

© . All rights reserved.