分布式训练(二): DeepSpeed 知识总结

发布时间:2026/5/31 21:35:05

分布式训练(二): DeepSpeed 知识总结 目录一、DeepSpeed 简介1.1 核心特性1.2 典型使用场景1.3 如果已经有 PP、TP为什么还需要 DeepSpeed1.3.1 PP / TP 解决的与未解决的1.3.2 DeepSpeed 在其中的作用1.3.3 关键结论1.4 混合精度训练详解1.4.1 不同张量的精度分工1.4.2 训练流程示意二、ZeRO 原理与阶段对比2.1 ZeRO 定义2.2 ZeRO-1分片优化器状态2.3 ZeRO-2分片优化器状态 梯度2.4 ZeRO-3分片优化器状态 梯度 参数2.5 显存对比以 7.5B 模型、64 卡为例2.6 扩展能力2.7 三个阶段优缺点与使用场景对比2.8 选择建议三、DeepSpeed 配置示例3.1 ZeRO-1 配置3.2 ZeRO-2 配置推荐中等规模模型3.3 ZeRO-3 配置大模型3.4 配置参数说明3.5 启动命令一、DeepSpeed 简介DeepSpeed 是微软开源的深度学习优化库主要用于大规模模型训练。1.1 核心特性ZeRO (Zero Redundancy Optimizer)将优化器状态、梯度、参数分片到多个 GPU大幅降低单卡显存占用可训练千亿/万亿参数模型3D 并行原生支持数据并行 (DP) 与流水线并行 (PP)张量并行 (TP) 主要通过与 Megatron-LM 集成Megatron-DeepSpeed实现从而组合出完整的 3D 并行混合精度训练(1.4 详细介绍)训练中同时使用低精度 (FP16/BF16) 和 FP32——前向/反向计算、激活和梯度用低精度以省显存、加速计算而 master 权重和优化器状态保留 FP32 以保证参数更新的数值精度CPU/NVMe Offload将参数和优化器状态卸载到 CPU 内存或 NVMe进一步突破显存限制DeepSpeed-MoE稀疏专家模型支持1.2 典型使用场景大模型训练训练 LLM、视觉大模型、MoE 模型等如 BLOOM、MT-NLG显存受限训练单卡/多卡显存不足时借助 ZeRO 与 Offload 训练更大模型框架集成可与 HuggingFace Transformers、PyTorch Lightning、Accelerate 等无缝配合1.3 如果已经有 PP、TP为什么还需要 DeepSpeedPP流水线并行和 TP张量并行只是「怎么切模型」的两种方法而 DeepSpeed 是「怎么训模型」的完整训练引擎。两者层次不同不是替代关系。1.3.1 PP / TP 解决的与未解决的PP / TP 解决模型切分按层切 / 按矩阵切让超大模型能放进多卡PP / TP 未解决DP 维度冗余、显存进一步压缩Offload、混合精度、训练工程细节、通信优化、分布式 checkpoint、并行调度策略等1.3.2 DeepSpeed 在其中的作用组件提供能力ZeRO在 DP 维度上分片优化器状态/梯度/参数消除冗余PP 引擎实现 1F1B 等流水线调度策略混合精度一行配置启用 BF16/FP16Offload把状态卸载到 CPU/NVMe突破显存限制Checkpoint自动处理分布式状态的保存与恢复通信优化overlap_comm、contiguous_gradients 等1.3.3 关键结论PP / TP 是切分策略DeepSpeed 是训练框架真实训练几乎都是TP PP DP 组合DP 维度需要 ZeRO 去冗余DeepSpeed 让用户只写 config不必手撸分布式细节1.4 混合精度训练详解混合精度指的是训练中同时使用 FP16/BF16 和 FP32低精度负责计算FP32 负责关键的累加和参数更新。1.4.1 不同张量的精度分工用途精度原因前向/反向计算FP16 / BF16计算快、显存省激活值 (activations)FP16 / BF16占大头省显存梯度 (gradients)FP16 / BF16通信快Master 权重FP32保证参数更新精度优化器状态 (m, v)FP32数值稳定性loss / 梯度累加FP32防止精度损失1.4.2 训练流程示意FP32 master weights │ 转 FP16/BF16 ▼ 低精度权重 ─► forward ─► loss ─► backward ─► 低精度梯度 │ 转 FP32 ▼ 优化器更新 (FP32) │ ▼ FP32 master weights更新补充为什么必须保留 FP32 主权重FP16 只有约 3-4 位有效数字若直接累加小的更新量如W1.0、lr*grad0.0001会被精度吞掉导致更新丢失所以必须在 FP32 上做累加再把更新后的权重转回低精度用于下一轮 forward。这里的FP32 主权重是训练时的内存策略与模型保存到磁盘的格式无关——训练 checkpoint 通常存 FP32 便于恢复训练而对外发布的权重如 LLaMA、Qwen大多是 BF16。二、ZeRO 原理与阶段对比2.1 ZeRO 定义ZeRO全称Zero Redundancy Optimizer零冗余优化器是 DeepSpeed 中的核心显存优化技术由微软在论文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》中提出。在传统的数据并行Data Parallelism, DP训练中每张 GPU 都会保存一份完整的模型参数、梯度和优化器状态存在大量冗余。ZeRO 的核心思想是将这些训练状态在数据并行的 GPU 之间进行分片partition每张卡只保存自己负责的那一部分从而消除冗余、大幅降低单卡显存占用使得在有限显存下能训练更大的模型。训练状态分为三类优化器状态 (Optimizer States)如 Adam 的 momentum、variance通常占用最多显存梯度 (Gradients)反向传播得到的梯度参数 (Parameters)模型权重本身根据分片的范围不同ZeRO 划分为三个递进的阶段ZeRO-1、ZeRO-2、ZeRO-3。2.2 ZeRO-1分片优化器状态分片内容仅优化器状态如 Adam 的 momentum、variance显存节省约 4 倍相比传统 DP通信开销与传统 DP 基本相同原理每张卡只保存 1/N 的优化器状态参数更新后通过 all-gather 同步2.3 ZeRO-2分片优化器状态 梯度分片内容优化器状态 梯度显存节省约 8 倍通信开销与传统 DP 基本相同reduce-scatter 替代 all-reduce原理每张卡只保留自己负责更新的那部分参数的梯度2.4 ZeRO-3分片优化器状态 梯度 参数分片内容全部三类状态显存节省与 GPU 数量 N 成线性关系理论上 N 倍通信开销约为传统 DP 的 1.5 倍原理参数也分片存储前向/反向计算时按需 all-gather用完即释放2.5 显存对比以 7.5B 模型、64 卡为例阶段单卡显存占用传统 DP~120 GBZeRO-1~31 GBZeRO-2~16 GBZeRO-3~1.9 GB2.6 扩展能力ZeRO-Offload在 ZeRO-2 基础上把优化器状态和梯度卸载到 CPUZeRO-Infinity在 ZeRO-3 基础上把参数卸载到 CPU/NVMe可训练万亿参数模型2.7 三个阶段优缺点与使用场景对比阶段优点缺点适用场景ZeRO-1- 实现简单对训练流程影响最小- 通信开销与传统 DP 相当- 训练速度快- 显存节省有限约 4 倍- 梯度和参数仍有冗余- 中小模型 1B 参数- 显存比较宽裕- 追求最快训练速度ZeRO-2- 显存节省较大约 8 倍- 通信开销仍与传统 DP 相当- 性价比高工业界使用最广- 参数仍未分片单卡需存完整模型权重- 模型大小受单卡显存限制- 中等规模模型1B ~ 10B 参数- 主流大模型微调场景- 显存与速度需平衡时的首选ZeRO-3- 显存节省最大随 GPU 数线性增长- 可训练超大模型百亿/千亿级- 可结合 Offload 突破显存上限- 通信开销增加约 50%- 训练速度变慢- 实现复杂调试较困难- 对网络带宽要求高- 超大模型 10B 参数训练- 单卡放不下完整模型时- 显存极度受限配合 Offload2.8 选择建议显存够用→ ZeRO-1 或 ZeRO-2通信少速度快显存不够→ ZeRO-3极端显存受限→ ZeRO-3 OffloadCPU/NVMe建议优先级ZeRO-2 → ZeRO-3 → ZeRO-3 Offload按需逐级升级避免过度牺牲速度三、DeepSpeed 配置示例3.1 ZeRO-1 配置{train_batch_size:32,optimizer:{type:AdamW,params:{lr:1e-4,weight_decay:0.01}},bf16:{enabled:true},zero_optimization:{stage:1}}3.2 ZeRO-2 配置推荐中等规模模型{train_batch_size:32,optimizer:{type:AdamW,params:{lr:1e-4,weight_decay:0.01}},bf16:{enabled:true},zero_optimization:{stage:2,overlap_comm:true,contiguous_gradients:true}}3.3 ZeRO-3 配置大模型{train_batch_size:32,optimizer:{type:AdamW,params:{lr:1e-4,weight_decay:0.01}},bf16:{enabled:true},zero_optimization:{stage:3,overlap_comm:true,contiguous_gradients:true,stage3_gather_16bit_weights_on_model_save:true}}3.4 配置参数说明顶层参数参数含义train_batch_size全局训练 batch size 单卡 batch × GPU 数 × 梯度累积步数DeepSpeed 会据此校验配置optimizer.type优化器类型常用AdamW、Adam、Lamb等optimizer.params.lr学习率optimizer.params.weight_decay权重衰减系数用于正则化防止过拟合bf16.enabled启用 BF16 混合精度训练。相比 FP16 数值范围更大、无需 loss scaling适合 A100/H100 等新硬件zero_optimization参数参数含义stageZeRO 阶段1分片优化器状态2加上梯度3再加上参数overlap_comm通信与计算重叠执行隐藏通信延迟、加速训练推荐开启contiguous_gradients将梯度复制到连续的显存块中减少反向传播时的内存碎片提升性能stage3_gather_16bit_weights_on_model_save仅 ZeRO-3 适用。保存模型时把分片的参数 all-gather 成完整权重再落盘便于后续加载3.5 启动命令deepspeed--num_gpus8train.py--deepspeed--deepspeed_configds_config.json

相关新闻