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更适合资源受限的环境，尤其在推理任务中，通过压缩存储和计算需求，提高了部署效率。合理运用这些格式能够优化性能和资源利用，推动大语言模型在更广泛场景中的应用。

2.2.2 显存(VRAM)需求计算 - 推理

以LLaMA 2 7B为例，显存需求如下：

仅考虑了模型参数本身，并未包括其他运行时所需的额外空间，如优化器状态、激活等。

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