DeepSpeed入門

深度學習模型正變得越來越複雜，訓練計算成本也在上升。由微軟開發的DeepSpeed能夠在最低資源下高效訓練大型模型。本章將引導你完成基本步驟，從安裝和環境設定到執行你的第一個模型，逐步學習如何使用DeepSpeed。

安裝DeepSpeed

在深入瞭解DeepSpeed的細節之前，我們首先需要安裝該庫。使用pip，這很容易實現：

pip install deepspeed

安裝過程中，你可能會看到類似下面的結果：

Collecting deepspeed
Downloading deepspeed-0.6.0-py3-none-any.whl (696 kB)
|████████████████████████████████| 696 kB 3.2 MB/s 
Collecting torch
Downloading torch-1.9.1-cp38-cp38-manylinux1_x86_64.whl (804.1 MB)
|████████████████████████████████|
deepspeed-0.6.0 torch-1.9.1 installed successfully

如果你願意，也可以克隆GitHub倉庫並從原始碼安裝：

git clone https://github.com/microsoft/DeepSpeed.git
cd DeepSpeed
pip install .

這將為你提供最新的功能，這些功能可能尚未在穩定版本中釋出。

環境設定

安裝DeepSpeed後，需要設定環境。首先，確保所有必需的依賴項都存在。

建立一個虛擬環境來管理依賴項：

python -m venv deepspeed-env
source deepspeed-env/bin/activate  # On Windows, use 'deepspeed-env\\Scripts\\activate'

如果你還沒有安裝PyTorch，請安裝：

pip install torch torchvision torchaudio

此外，根據你的用例，你可能需要CUDA或其他型別的GPU加速。如果你使用的是裝有GPU的機器，安裝CUDA版本的PyTorch就像在終端執行以下命令一樣簡單：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

這將確保DeepSpeed利用你機器的所有硬體能力。

基本概念和術語

在執行你的第一個模型之前，讓我們瞭解一些你在DeepSpeed中經常會遇到的基本概念和術語。

最佳化器 - DeepSpeed目前支援多種最佳化器，可用於最佳化大型模型的訓練。最佳化器在訓練模型時處理梯度更新。
排程器 - 排程器在訓練期間更新學習率。預設情況下，DeepSpeed集成了所有PyTorch排程器，並進一步提供了為大型模型開發的其他自定義排程器。
零冗餘最佳化器 (ZeRO) - 這是一種記憶體最佳化技術，透過將模型狀態劃分為多個GPU來減少大型模型的記憶體佔用。
累積梯度 - 透過在模型權重更新之前對多個迭代的梯度求和，這可以促進使用比GPU記憶體允許的更大的批次大小。
檢查點啟用 - 這以增加計算為代價節省了一些記憶體，在反向傳播期間重新計算前向傳遞啟用。

理解這些概念應該足以讓你理解DeepSpeed的大部分高階功能並自定義你的訓練流程。

使用DeepSpeed執行你的第一個模型

現在你的環境已經設定好，並且你已經熟悉了基本術語，讓我們執行一個簡單的DeepSpeed模型。我們首先建立一個基本的PyTorch模型，然後新增DeepSpeed來檢視效能提升。

步驟1：建立一個簡單的PyTorch模型

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(10, 50) # input layer (10) -> hidden layer (50)
self.fc2 = nn.Linear(50, 1) # hidden layer (50) -> output layer (1)

def forward(self, x):
x = torch.relu(self.fc1(x)) # hidden layer activation function
x = self.fc2(x)
return x

model = SimpleModel()

步驟2：實現DeepSpeed

現在，讓我們重構程式碼使其與DeepSpeed一起工作。我們將使用DeepSpeed和一些基本配置初始化模型。

import deepspeed

ds_config = {
   "train_batch_size": 32,
   "fp16": {
      "enabled": True
   },
   "zero_optimization": {
      "stage": 1
   }
}

model_engine, optimizer, _, _ = deepspeed.initialize(
   model=model,
   model_parameters=model.parameters(),
   config=ds_config
)

輸出

如果一切順利，DeepSpeed將初始化並打印出配置設定：

[INFO] DeepSpeed info: version=0.6.0, git-hash=unknown, git-branch=unknown
[INFO] Initializing model parallel group with size 1
[INFO] Initializing optimizer with DeepSpeed Zero Optimizer

步驟3：訓練模型

此時，你應該能夠使用DeepSpeed訓練你的模型。下面是一個示例訓練迴圈。

for epoch in range(5) − 
inputs = torch.randn(32, 10)
labels = torch.randn(32, 1)

model_engine.train()
outputs = model_engine(inputs)
loss = nn.MSELoss()(outputs, labels)

model_engine.backward(loss)
model_engine.step()
print(f'Epoch {epoch+1}, Loss: {loss.item()}')

輸出

每個epoch都會給你類似這樣的結果

Epoch 1, Loss: 0.4857
Epoch 2, Loss: 0.3598
Epoch 3, Loss: 0.2893
Epoch 4, Loss: 0.2194
Epoch 5, Loss: 0.1745

步驟4：儲存模型

最後，你可以儲存到目前為止訓練的模型：

model_engine.save_checkpoint('./checkpoint', epoch=5)

輸出

[INFO] Saving model checkpoint to ./checkpoint

DeepSpeed的高階功能

在對DeepSpeed有了基本的瞭解之後，讓我們看看DeepSpeed的一些高階功能。這些高階功能是為了處理訓練大型模型的複雜性，減少記憶體消耗並提高計算效率而實現的。

混合精度訓練FP16 - DeepSpeed快速模型訓練的原因之一是它支援透過使用半精度進行混合精度訓練。
ZeRO最佳化階段 - DeepSpeed擁有一種名為ZeRO的改變遊戲規則的技術，它透過將模型狀態劃分為多個GPU來減少記憶體。
梯度累積 - DeepSpeed支援的另一種策略是梯度累積，它可以在不需要更多GPU記憶體的情況下模擬更大的批次大小。
解除安裝 - 即使對於非常大的模型，ZeRO Stage 3提供的最佳化也可能不夠。

總結

DeepSpeed入門的關鍵步驟包括安裝庫、設定環境、瞭解一些基本概念以及執行你的第一個模型。DeepSpeed允許以更高的效率、更高的記憶體和更低的整體訓練時間來訓練大型模型。本章將使你能夠進一步深入瞭解DeepSpeed的高階功能，從而推動你的深度學習專案。

列印頁面