Distributed Training: DeepSpeed ZeRO 1/2/3 + Accelerate, Megatron-LM

ForHHeart發表於2024-09-07
AI

1 Introduction

Github: https://github.com/microsoft/DeepSpeed

  1. ZeRO: Memory Optimizations Toward Training Trillion
    Parameter Models
  2. ZeRO-Offload: Democratizing Billion-Scale Model Training
  3. ZeRO-Infinity: Breaking the GPU Memory Wall
    for Extreme Scale Deep Learning
  4. ZeRO++: Extremely Efficient Collective Communication for Giant
    Model Training

ZeRO(Zero Redundancy Optimizer)是一種去除冗餘的分散式資料並行(Data Parallel)方案，分為Stage 1, Stage 2, Stage 3，而Deepspeed就是論文中ZeRO方法的Microsoft官方的工程實現。

ZeRO-Offload為解決由於ZeRO而增加通訊資料量的問題，提出將GPU轉移到CPU

ZeRO-Infinity同樣是進行offload，ZeRO-Offload更側重單卡場景，而ZeRO-Infinity則是典型的工業界風格，奔著極大規模訓練去了

ZeRO++是對ZeRO 3的通訊最佳化，最佳化了以下三個方面：

  1. 每個伺服器有完整的模型引數，消除跨伺服器的All_gather操作；
  2. 通訊時，基於塊的量化，模型引數從FP16轉換成INT8；
  3. 替代ring-based ReduceScatter通訊，改為分層級的量化 AllToALL；

Megatron-LM是NVIDIA開發的大規模語言模型訓練框架，相比於DeepSpeed而言，具有更好的模型並行和流水線並行技術，但資料並行DeepSpeed更有優勢。

2 預備知識

2.1 分散式並行策略

Tensor Parallel, Data Parallel, Model Parallel, Pipeline Parallel

2.2 LLM推理和訓練的算力需求估算

2.2.1 資料精度格式

對於大語言模型，選擇合適的精度格式至關重要。高精度如FP32適合高要求的任務，但消耗資源多；FP16和bfloat16則在維持效能的同時，顯著降低了計算成本。低精度格式如int8和fp4更適合資源受限的環境，尤其在推理任務中，透過壓縮儲存和計算需求，提高了部署效率。合理運用這些格式能夠最佳化效能和資源利用，推動大語言模型在更廣泛場景中的應用。

名稱 簡稱 對應位元組 對應位元
單精度浮點格式(Single-precision floating-point format) fp32 4 Bytes 32 bits
半精度浮點格式(Half-precision floating-point format) fp16 2 Bytes 16 bits
腦浮點格式(Brain floating-point format) bp16 2 Bytes 16 bits
8位整數格式(8-bit integer format) int8 1 Bytes 8 bits
4位浮點格式(4-bit floating-point format) fp4 0.5 Bytes 4 bits
4位正常浮點格式 (4-bit NormalFloat format) nf4 0.5 Bytes 4 bits

2.2.2 視訊記憶體(VRAM)需求計算 - 推理

以LLaMA 2 7B為例，視訊記憶體需求如下：

僅考慮了模型引數本身，並未包括其他執行時所需的額外空間，如最佳化器狀態、啟用等。

型別 模型精度 模型規模 推理/訓練 最低視訊記憶體(以粗略計算方式)
全精度 FP32 7B 推理 7B * 4 Bytes = 28 GB
半精度 FP16 7B 推理 7B * 2 Bytes = 14 GB
低精度 INT8 7B 推理 7B * 1 Bytes = 7 GB
INT4 7B 推理 7B * 0.5 Bytes = 3.5 GB

3 ZeRO

Distributed Training: DeepSpeed ZeRO 1/2/3 + Accelerate, Megatron-LM

如今，Mixed-Precision Training和Adam Optimizer是LLM Distributed Training的標配

ZeRO將模型訓練階段，每張卡中視訊記憶體內容分為兩類(以最佳化器Adam為例)：

  1. Model States:
    • Parameters(fp16)
    • Gradient(fp16)
    • Optimizer States(fp32)

VRAM計算：假設Parameters是\(\Psi\)，則共需要2\(\Psi\)+2\(\Psi\)+(4\(\Psi\)+4\(\Psi\)+4\(\Psi\))=16\(\Psi\) bytes進行儲存。

  1. Residual States
    • 啟用值activation
    • 臨時緩衝區buffer
    • 無法使用的視訊記憶體碎片fragmentation

3.1 Stage 1, 2, 3

Distributed Training: DeepSpeed ZeRO 1/2/3 + Accelerate, Megatron-LM
  • \(P_{os}\)是指

3.2 ZeRO-Offload

3.2.1 通訊資料量分析

3.3 ZeRO-Infinity

3.4 ZeRO++

4 DeepSpeed + Accelerate

pip install deepspeed
accelerate config

image

5 Megatron-LM

Github: https://github.com/NVIDIA/Megatron-LM

  1. Megatron-LM: Training Multi-Billion Parameter Language Models Using
    Model Parallelism
  2. Efficient Large-Scale Language Model Training on GPU Clusters
    Using Megatron-LM
  3. Reducing Activation Recomputation in Large Transformer Models

Reference

  • ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters | Microsoft Research Blog
  • DeepSpeed之ZeRO系列：將視訊記憶體最佳化進行到底 | 知乎
  • 分散式訓練框架Megatron-LM程式碼概覽 | Bilibili
  • Zero++分散式並行 資料並行 | Bilibili

相關文章