使用 DeepSpeed 進行記憶體最佳化

隨著深度學習模型複雜性和大規模計算的不斷增長，記憶體最佳化在訓練過程中至關重要。DeepSpeed 提供了多種節省記憶體的技術，例如解除安裝、梯度檢查點和 ZeRO。開發人員可以基於這些節省記憶體的技術，在標準硬體上訓練非常龐大的模型。此外，這些技術還可以訓練以前受硬體限制的模型。

DeepSpeed 不僅在研究領域取得了令人矚目的進步，在工業領域也同樣取得了成功，因此已成為深度學習從業者不可或缺的工具。這意味著透過這些策略，您可以使模型消耗更少的記憶體，並真正突破硬體的限制。

為什麼需要記憶體最佳化？

記憶體最佳化是訓練深度學習模型中最關鍵的組成部分之一。對於像 GPT 和 BERT 這樣擁有數十億引數的模型，在可用硬體上進行訓練時，必須有效地管理記憶體。DeepSpeed 是一個開源庫，用於訓練深度學習模型，它具有 ZeRO 最佳化器、解除安裝技術和梯度檢查點等功能，以避免訓練過程中的主要記憶體佔用。

深度學習中的記憶體問題

深度學習模型的規模和複雜性近年來顯著增加。這些大型模型需要大量的記憶體進行訓練。微軟開發的深度學習最佳化庫 DeepSpeed 為這些挑戰提供了強大的解決方案。

深度學習模型有點像藝術，它們代表著新興事物，隨著模型規模的增長，記憶體相關的問題也會隨之出現。一些最常見的記憶體相關問題包括：

模型引數 - 大型模型，如 GPT-3，擁有數千億個引數，因此需要大量的記憶體來儲存。
梯度 - 在訓練過程中計算每個引數的梯度也必須計算並儲存在記憶體中，這會消耗更多的記憶體。
啟用對映 - 前向傳遞過程中產生的所有中間值都需要儲存，直到反向傳遞僅用於梯度計算，這稱為啟用對映。
批大小 - 更大的批大小可以提高收斂速度，但會消耗更多的記憶體。
資料並行 - 將資料在多個 GPU 之間進行分發是減少訓練時間的好策略，但毫無疑問，它確實會消耗大量的記憶體，除非得到控制。

除非識別出這些陷阱，否則即使在消費級硬體上訓練大型模型也變得不可能。DeepSpeed 透過使用創新的節省記憶體技術克服了這些挑戰。

DeepSpeed 的記憶體最佳化技術

DeepSpeed 在訓練模型時有多種方法可以最佳化記憶體使用。一些方法包括 ZeRO（零冗餘最佳化器）、梯度檢查點和啟用重計算。

零冗餘最佳化器 (ZeRO)
梯度檢查點
解除安裝技術

1. 零冗餘最佳化器 (ZeRO)

ZeRO 主要關注在最佳化器狀態、梯度和模型引數的冗餘副本被移除的地方進行記憶體最佳化。ZeRO 經歷以下三個階段：

階段 1 - 將最佳化器狀態跨 GPU 分片，每個 GPU 儲存一部分最佳化器狀態。
階段 2 - 進一步減少記憶體，因為梯度跨 GPU 分片。
階段 3 - 模型引數被分片，現在可以訓練高達萬億引數的模型。

示例

import deepspeed

model = MyModel()  # your dl model
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# DeepSpeed configuration for ZeRO
ds_config = {
 "train_batch_size": 8,
 "zero_optimization": {
    "stage": 2,   # adjust the stage of ZeRO here
    "allgather_partitions": True,
    "reduce_scatter": True,
    "allgather_bucket_size": 5e8,
    "overlap_comm": True,
    "contiguous_gradients": True
 }
}
# Initialize DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
  model=model,
  optimizer=optimizer,
  config_params=ds_config
)

# Training loop
for batch in train_dataloader:
    loss = model_engine(batch)
model_engine.backward(loss)
model_engine.step()

您會注意到記憶體使用量大大降低，特別是對於大型模型。然後，記憶體分析器可以突出顯示 ZeRO 最佳化啟動的位置。

2. 梯度檢查點

梯度檢查點透過不在緩衝區中儲存前向傳遞期間的啟用來減少記憶體。相反，它們在反向傳遞中被重建，犧牲一點計算來節省一些記憶體。

示例

import torch
from torch.utils.checkpoint import checkpoint

def custom_forward(*inputs):
    return model(*inputs)

# Gradient checkpointing
outputs = checkpoint(custom_forward, input_data)
loss = criterion(outputs, labels)
loss.backward()

在這種情況下，節省的記憶體將取決於中間啟用的大小。

3. 解除安裝技術

DeepSpeed 還提供另一種形式的解除安裝。它允許您將模型的部分內容（如最佳化器狀態和梯度）移動到 CPU 甚至 NVMe 儲存器，從而釋放 GPU 記憶體以供其他用途使用。

CPU 解除安裝

DeepSpeed 允許我們將最佳化器狀態和梯度解除安裝到 CPU。這也釋放了寶貴的 GPU 記憶體。如果 GPU 上的記憶體有限，但 CPU 上的記憶體相當多，這將非常有用。

示例

ds_config = {
    "train_batch_size": 8,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        }
    }
}
model_engine, optimizer, _, _ = deepspeed.initialize(
   model=model,
   optimizer=optimizer,
   config_params=ds_config
)

由於解除安裝到 CPU 的傳輸涉及裝置間通訊成本，因此訓練速度相對較慢，但在那些原本無法適應記憶體受限 GPU 的模型大小上，它仍然很有用。

NVMe 解除安裝

對於大型模型，這還不夠。DeepSpeed 還將最佳化器狀態和梯度解除安裝到 NVMe 儲存器。這將進一步提高模型的規模，即使在訓練中也能實現。

示例

ds_config = {
    "train_batch_size": 8,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/local_nvme"
        },
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/local_nvme"
        }
    }
}
model_engine, optimizer, _, _ = deepspeed.init(
   model=model,
   optimizer=optimizer,
   config_params=ds_config
)

使用 NVMe 進行解除安裝將能夠訓練海量模型，儘管速度將在很大程度上取決於 NVMe 驅動器的 I/O 速度。

記憶體最佳化案例研究

讓我們討論一些使用 DeepSpeed 進行記憶體最佳化的真實案例研究 -

案例研究 1：使用 ZeRO 最佳化訓練 GPT-2

使用 DeepSpeed，一個研究團隊將這個擁有 15 億個引數的 GPT-2 模型的訓練擴充套件到了消費級 GPU。使用 ZeRO 階段 3，可以在 4 個 NVIDIA RTX 3090 GPU 上對其進行訓練，每個 GPU 的總記憶體為 24 GB。如果沒有使用 ZeRO，訓練將是不可能的，因為該模型每個 GPU 需要超過 50 GB 的記憶體。

案例研究 2：使用 NVMe 解除安裝 1750 億引數模型

微軟利用 DeepSpeed 的解除安裝功能，在一個記憶體有限的 GPU 叢集上訓練了一個 1750 億引數的模型。在模型訓練期間，幾乎沒有記憶體瓶頸來解除安裝最佳化器狀態和引數，這表明即使 GPU 資源有限，解除安裝也可以為超大型模型鋪平道路。

列印頁面