一句话定义:DeepSpeed 是微软开源的一套分布式训练优化库,主打“省显存”;Megatron(含 Megatron-LM 和 Megatron-Core)是英伟达开源的大模型训练框架,主打“切模型”。两者都在解决同一个问题:一张显卡装不下的大模型,怎么拆到成百上千张卡上训练。
先看省显存这条线。 最朴素的训练方式是数据并行(data parallelism,DP):每张卡都存一份完整的模型参数、梯度和优化器状态,只是喂不同的数据。这在模型小的时候没问题;但优化器状态(例如 Adam 的一阶、二阶动量)往往比参数本身还占空间。DeepSpeed 的核心贡献叫 ZeRO(Zero Redundancy Optimizer,零冗余优化器),思路很直白:既然每张卡都存一份一模一样的副本太浪费,那就切开分给各卡保管——第一层切优化器状态,第二层连梯度一起切,第三层连参数也切;需要用时再用通信把片段临时拼回来,算完再扔掉。此外,DeepSpeed 还提供把状态卸载到 CPU 内存乃至本地磁盘的选项、激活重计算等配套手段。
打个比方:一个装修队要给十套一模一样的房子施工。数据并行是每个工人都背一整套图纸和全套工具,各干一套房;ZeRO 是把图纸和工具拆开,每人只背一小包,干到哪一步就朝其他人借哪一页。背包轻了,能接的活儿就大了,代价是彼此喊话更频繁,也就是通信开销。
再看切模型这条线。 Megatron 走的是模型并行(model parallelism):如果单层的一个矩阵乘法大到一张卡放不下,就把这一层按行或列切开,几张卡同时算再拼起来,这叫张量并行(tensor parallelism,TP);如果模型层数太多,就按层切成若干段,第一段算完把中间结果传给第二段,这叫流水并行(pipeline parallelism,PP),代价是流水线开头和结尾会有工人闲着,也就是常说的气泡(bubble)。
和 FSDP、ZeRO 的关系,容易混,列个表:
| 名字 | 切的是什么 | 一句话定位 |
|---|---|---|
| DeepSpeed / ZeRO | 数据并行副本里的优化器状态、梯度、参数 | 消除冗余副本,省显存 |
| FSDP | 参数、梯度、优化器状态全切分 | PyTorch 原生实现的同类思路 |
| Megatron TP | 单层内部的矩阵运算 | 一层都放不下一张卡时用 |
| Megatron PP | 模型按层分段 | 跨机器拼出超大模型 |
关键点是:ZeRO 是一种“想法”,DeepSpeed 是把它实现出来的引擎;FSDP(Fully Sharded Data Parallel)是 PyTorch 官方按同一思路做的全分片方案,和 ZeRO 的第三层属于同一档,通常二选一即可,不必同时开。而 Megatron 切的是模型结构本身,和 ZeRO、FSDP 切的是“副本冗余”,两者方向正交,可以叠加使用——常见的超大模型训练就是 TP 加 PP 把模型摊开,再用 ZeRO 的第一层或 FSDP 消掉数据并行组里的冗余。具体支持哪些组合、哪些开关,以各自官方文档为准。
对从业者的实际意义:这是一份选型顺序。单机多卡、模型放得下,普通 DDP 就够;显存吃紧,先上激活重计算和 ZeRO 前两层;还不够,再考虑 ZeRO 第三层或 FSDP;单层就放不下,才需要 TP(尽量放在同一台机器内,因为通信最密集);跨节点拼规模用 PP。对普通职场人来说,值得知道的是:大模型训练之所以从少数实验室的“炼金术”变成可复制的工程,靠的正是这类把显存和通信算到字节的框架。
