收藏必学!大模型分布式训练核心:ZeRO技术详解与实战应用!
本文深入解析分布式训练中的ZeRO技术,详细介绍了ZeRO-1、ZeRO-2和ZeRO-3的原理、区别与应用场景。ZeRO通过将模型参数、梯度和优化器状态分片到不同GPU,显著降低显存占用,从而支持千亿甚至万亿参数级别的大模型训练。文章还对比了不同ZeRO版本的适用场景,帮助开发者根据实际需求选择最优的分布式训练策略。
在大模型训练中,显存瓶颈是制约模型规模的重要因素。DeepSpeed 推出的Zero Redundancy Optimizer(ZeRO) 系列,通过将模型状态(参数、梯度、优化器状态)分布到不同 GPU 上,实现显著的显存优化。

本文将详细介绍 ZeRO-1、ZeRO-2、ZeRO-3 的原理、区别与应用场景。
希望大家带着下面的问题来学习,我会在文末给出答案。
1.ZeRO 系列优化器的核心原理是什么?
2.ZeRO-1、ZeRO-2、ZeRO-3 有哪些区别?
3.不同 ZeRO 版本适合的应用场景是什么?
一、ZeRO 核心原理
ZeRO的全称为 Zero Redundancy Optimizer,意思是去除冗余的优化器,我们都知道在分布式训练中,主要的参数包括三部分:模型参数(Parameters)、优化器状态(Optimizer States)、梯度(Gradients),其中优化器状态会占据大约2倍参数量的显存空间,ZeRO 的目标就是 减少每个 GPU 的显存占用,从而支持更大模型训练。它通过 分布式存储模型状态 来降低冗余:
·传统数据并行问题
o每个 GPU 都保存完整的参数、梯度和优化器状态。
o对于千亿参数模型,这会导致显存占用过高,限制可训练模型规模。
·ZeRO 优化思路
o将模型状态按 GPU 分片(Shard),每个 GPU 只存自己负责的一部分。
o在前向和反向传播计算时,通过通信获取必要的参数和梯度,实现显存按需使用。

二、 ZeRO-1:优化器状态分片
ZeRO-1仅将 优化器状态(如 Adam 的动量、平方梯度)分片到不同 GPU,参数和梯度仍是完整拷贝。
模型训练过程中正向传播和反向传播中并不会用到优化器状态,只有在梯度更新的时候才会使用梯度和优化器状态计算新参数,因此每个GPU单独使用一段优化器状态,对各自GPU的参数更新完之后,再把各个GPU的模型参数合并成完成的模型。
它的训练过程与DDP类似,forward过程由每个rank的GPU独自完整的完成,然后进行backward过程,在backward过程中,梯度通过AllReduce进行同步。
ZeRO-1使显存占用下降,支持更大 batch-size 或更大模型,并且实现简单,通信开销低,适用于模型参数较大,但优化器状态占显存比例高的场景。
三、ZeRO-2:优化器状态 + 梯度分片
ZeRO第二阶段在 ZeRO-1 基础上,进一步将 梯度 分片。前向计算后,GPU 只保存本地梯度的部分,反向传播完成后通过 AllReduce 聚合梯度。
ZeRO显存节省比 ZeRO-1 更大,支持更大模型和更高 batch-size。AllReduce 用于聚合梯度,需要一定通信开销,但与显存节省相比收益明显。适用于模型参数大、梯度占用显存较多的场景,在多节点训练中尤为有效。
四、 ZeRO-3:全模型状态分片
ZeRO-3将 参数、梯度、优化器状态 全部分片到不同 GPU。每个 GPU 只持有自己负责的模型片段,计算时通过通信获取必要片段。
ZeRO-3显存占用最低,理论上可以训练无限大模型(受通信和计算限制),支持千亿、万亿参数级别模型训练。参数和梯度都需要按需通信,通信开销较 ZeRO-1/2 高,需要精心设计分片策略和通信调度以提升效率。
适用于超大模型训练(千亿参数以上),多节点训练、大 batch-size、混合精度训练场景。
五、图解ZeRO
官方给出了一个5分钟的解释视频,我们一帧帧的看一下:
- 我们拥有一个由 16 个 Transformer 块组成的模型,每个块都是一个完整的 Transformer 模块。

- 在训练中,我们使用一个大型数据集并借助 4 张 GPU 进行训练。

- 为此,采用了 三阶段分布式策略,对模型参数(Optimizer + Parameter + Gradient,简称 OPG)和数据同时进行切分与分配,使得它们分布在四张 GPU 上。

4.每个模块下的方格表示该模块的显存占用情况:
·第一行:表示模型的 FP16 权重参数;
·第二行:是用于反向传播的 FP16 梯度,用以更新权重;
·其余大部分绿色区域对应的是 优化器状态所占显存,包括 FP32 精度的梯度、副方差、动量以及参数。这些优化器状态在计算完 FP16 梯度后才会被加载使用。由于 ZeRO Stage 3 采用混合精度训练,因此在前向传播中模型参数均以半精度(FP16)进行计算。
此外,每个模块还需要分配一部分显存用于保存激活值(Activation),图中以蓝色部分表示。

5.在ZeRO3模式下,每张GPU仅负责模型的一部分参数,训练流程如下:
(1)参数广播与前向计算
·首先,GPU_0 将自己持有的模型分块权重 broadcast 给其他 GPU。

·所有 GPU 收到该权重后,除 GPU_0 外,其余 GPU 会将其暂存于缓存中。
·各 GPU 使用这部分参数在各自的数据上执行前向传播,并保存相应层的激活值。

·当前层的计算完成后,临时参数从除 GPU_0 外的所有设备中释放。

(2)依次广播其余分块
·接着,GPU_1 将其负责的参数广播给所有 GPU,并重复上述前向计算与释放过程。

·然后是 GPU_2、GPU_3 依次进行广播与计算,直到全部 16 个 Transformer 块的前向传播完成。此时,每张 GPU 都基于各自的数据计算出一个损失值。

(3)反向传播阶段
·反向传播从最后一个模型分块(假设在 GPU_3 上)开始。所有 GPU 载入该分块权重,利用之前保留的激活值计算梯度。
·其他 GPU 会将本地计算得到的该分块梯度发送至 GPU_3,用于梯度聚合与累积。最终,GPU_3 完成权重更新并保存更新后的参数。

·各 GPU 清理相应的临时权重、梯度和激活缓存。
(4)继续回传并更新前面的分块
·GPU_2 将其模型参数广播出去,其他 GPU 利用这些参数继续执行反向传播并计算梯度。
·重复该过程,直至每个 GPU 持有的分块都完成了梯度计算与参数更新。

(5)优化器更新与精度回写
·在所有 GPU 上并行启动优化器更新操作。
·优化器会先将模型权重转换为 FP32 精度执行权重更新,然后再转换回 FP16 精度,用于下一次迭代的前向传播。

至此,一个完整的 ZeRO3 分布式训练迭代完成。
最后,我们回答文章开头提出的问题
1.ZeRO 系列优化器的核心原理是什么?
通过将模型状态(参数、梯度、优化器状态)分片到不同 GPU,显存按需使用,显著降低每个 GPU 的显存占用,从而支持大模型训练。
2.ZeRO-1、ZeRO-2、ZeRO-3 有哪些区别?
·ZeRO-1:只分片优化器状态。
·ZeRO-2:分片优化器状态 + 梯度。
·ZeRO-3:分片优化器状态 + 梯度 + 参数。显存节省和通信开销随版本递增。
3.不同 ZeRO 版本适合的应用场景是什么?
·ZeRO-1:优化器状态占显存高的中大型模型训练。
·ZeRO-2:模型和梯度占显存高,多节点训练场景。
·ZeRO-3:超大模型(千亿参数及以上)、大 batch-size、多节点训练。
六、 AI大模型从0到精通全套学习大礼包
我在一线互联网企业工作十余年里,指导过不少同行后辈。帮助很多人得到了学习和成长。
只要你是真心想学AI大模型,我这份资料就可以无偿共享给你学习。大模型行业确实也需要更多的有志之士加入进来,我也真心希望帮助大家学好这门技术,如果日后有什么学习上的问题,欢迎找我交流,有技术上面的问题,我是很愿意去帮助大家的!
如果你也想通过学大模型技术去帮助就业和转行,可以扫描下方链接👇👇
大模型重磅福利:入门进阶全套104G学习资源包免费分享!

01.从入门到精通的全套视频教程
包含提示词工程、RAG、Agent等技术点
02.AI大模型学习路线图(还有视频解说)
全过程AI大模型学习路线


03.学习电子书籍和技术文档
市面上的大模型书籍确实太多了,这些是我精选出来的


04.大模型面试题目详解


05.这些资料真的有用吗?
这份资料由我和鲁为民博士共同整理,鲁为民博士先后获得了北京清华大学学士和美国加州理工学院博士学位,在包括IEEE Transactions等学术期刊和诸多国际会议上发表了超过50篇学术论文、取得了多项美国和中国发明专利,同时还斩获了吴文俊人工智能科学技术奖。目前我正在和鲁博士共同进行人工智能的研究。
所有的视频由智泊AI老师录制,且资料与智泊AI共享,相互补充。这份学习大礼包应该算是现在最全面的大模型学习资料了。
资料内容涵盖了从入门到进阶的各类视频教程和实战项目,无论你是小白还是有些技术基础的,这份资料都绝对能帮助你提升薪资待遇,转行大模型岗位。


智泊AI始终秉持着“让每个人平等享受到优质教育资源”的育人理念,通过动态追踪大模型开发、数据标注伦理等前沿技术趋势,构建起"前沿课程+智能实训+精准就业"的高效培养体系。
课堂上不光教理论,还带着学员做了十多个真实项目。学员要亲自上手搞数据清洗、模型调优这些硬核操作,把课本知识变成真本事!

如果说你是以下人群中的其中一类,都可以来智泊AI学习人工智能,找到高薪工作,一次小小的“投资”换来的是终身受益!
应届毕业生:无工作经验但想要系统学习AI大模型技术,期待通过实战项目掌握核心技术。
零基础转型:非技术背景但关注AI应用场景,计划通过低代码工具实现“AI+行业”跨界。
业务赋能 突破瓶颈:传统开发者(Java/前端等)学习Transformer架构与LangChain框架,向AI全栈工程师转型。
👉获取方式:
😝有需要的小伙伴,可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓
更多推荐



所有评论(0)