DeepSpeed 介紹

DeepSpeed 是什麼？

DeepSpeed 是 Microsoft 開源的 分布式深度學習訓練優化庫，專門幫你用少量 GPU 訓練超大模型（數十億、甚至千億參數），核心在於高效的 分布式資料、模型和優化器並行化。

主要用途

• 需要訓練大語言模型（GPT、T5）、超大視覺模型。
• 想把 GPU 記憶體壓到極限，以少量硬體做最大化模型規模。
• 要在多節點、多 GPU 上做分布式深度學習。

核心技術：ZeRO (Zero Redundancy Optimizer)

• 目的是解決大模型「單 GPU 記憶體不足」的問題。
• 將 optimizer state、梯度、參數 分散到多張 GPU。
• ZeRO 分三個階段（Stage 1、2、3）：
  • Stage 1：分散 optimizer states
  • Stage 2：再分散梯度
  • Stage 3：連參數本身都分散，記憶體需求最低
• 例如 ZeRO Stage 3 可以將模型顯存需求降低 10 倍以上。

支援的並行類型

• Data Parallelism（數據並行）：同一模型在不同 GPU 各自算不同 mini-batch。
• Model Parallelism（模型並行）：不同 GPU 負責模型不同層。
• Pipeline Parallelism（管線並行）：不同 GPU 依序負責模型不同階段。
• DeepSpeed 可以和 Megatron-LM 結合，做 tensor-level parallelism。

常用功能

• 混合精度（FP16 / BF16），自動管理 loss scaling
• Gradient checkpointing，降低顯存消耗
• 高效推理：DeepSpeed Inference Engine
• 適用 PyTorch，與 Hugging Face Transformers 整合方便

安裝方法

使用 PIP 或 要支援特定 CUDA/NCCL，可使用官方安裝指令

pip install deepspeed

DS_BUILD_OPS=1 DS_BUILD_AIO=1 DS_BUILD_UTILS=1 pip install .

常見的工程用法

將 PyTorch 代碼的 optimizer 和模型包進 deepspeed.initialize()，並透過 JSON 配置檔（ds_config.json）控制並行、ZeRO、混合精度、學習率調整等。

import deepspeed

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

for batch in dataloader:
    outputs = model_engine(batch)
    loss = loss_fn(outputs, labels)
    model_engine.backward(loss)
    model_engine.step()