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

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

如今，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進行儲存。

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

3.1 Stage 1, 2, 3

• \(P_{os}\)是指

3.2 ZeRO-Offload

3.2.1 通訊資料量分析

3.3 ZeRO-Infinity

3.4 ZeRO++

4 DeepSpeed + Accelerate

pip install deepspeed
accelerate config

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