1. 为什么大模型训练绕不开这三套方案1.1 显存危机与三种并行的底层逻辑先聊一个最直接的问题为什么要把模型切碎来训练假设你手上是一个70B参数的模型用FP16训练。参数本身占70B × 2Byte 140GBAdam优化器需要额外保存一阶动量fp32和一阶动量对应的fp32参数副本加起来约140GB × 4Byte × 2 280GB再加上梯度约140GB。这一口气就是560GB朝上。而一张主流训练卡只有80GB显存单卡根本装不下更别说算起来卡得没法看。所以分布式训练解决的不是“变快”而是“装得下”和“算得动”。传统的数据并行Data Parallel只是把同一份模型复制到每张卡上每张卡算不同的batch最后同步梯度。这个方案在模型小的时候很好用但模型一大每张卡都要完整放下一份模型和优化器状态显存瓶颈立刻出现。于是就有了三种更精细的并行思路张量并行Tensor Parallel把一个层内部的矩阵拆开让多张卡合算一个层。典型代表就是Megatron-LM的列并行、行并行。流水线并行Pipeline Parallel把若干层切成几段每张卡负责一段数据像流水线一样流过各段。参数分片Parameter Sharding不切计算只切存储。每张卡只保存模型的一部分参数、梯度和优化器状态计算时再把需要的参数临时收齐。FSDP2和DeepSpeed ZeRO-3走的就是这条路。前两种切法本质是“计算并行”必然带来矩阵切分和通信开销第三种切法本质是“存储并行”通信开销集中在参数的全收集all-gather和梯度的归纳reduce-scatter上。三者的代价形态完全不同这也是选型时要权衡的核心。1.2 三套方案各自的定位FSDP2是PyTorch原生生态在2.x版本后主推的全参数分片方案站在用户视角它几乎不需要改模型结构只要层层标记分片点就能把一个大模型塞进几十张卡里跑起来。DeepSpeed ZeRO是另一套被广泛使用的参数分片框架它把优化器状态、梯度、参数三个对象按阶段切开ZeRO-3达到的效果和FSDP在概念上等价但实现上和工程生态有明显差异。Megatron-LM是面向超大规模模型的并行框架强调的是张量并行和流水线并行配合数据并行和梯度累积擅长把几百张卡以上的训练做得又快又稳。它不以“参数分片”为主打而是以“把计算切碎”为核心。我见过不少团队在踩坑后才明白这三个方案不是同一件事的三种替代品而更像是“存储分片”和“计算分片”两个维度的不同组合。下面我把每个方案的核心原理拆开来讲再给出一份可以直接套用的选型思路。2. FSDP2核心原理与实现细节2.1 从FSDP1到FSDP2接口到底改了什么第一代FSDP在一开始确实解决了“模型太大装不下”的问题但用过的人大概都有感触它的实现方式是先将一个模块的所有参数拍扁成一个“FlatParameter”再对这个大参数做分片。这样做的缺点是参数在显存里必须整块连续存放遇到某些模型结构时反而多出一份内存拷贝同时所有参数的收集、释放都绑定在模块级别的钩子上调度上不够精细。FSDP2不再使用FlatParameter而是把每个参数单独分片并且用DTensor来描述“这个参数在哪些设备上以什么方式分布”。从上层看每个参数仍然是一个普通的nn.Parameter但它的底层存储被拆到多张卡上。优化器在参数分片的基础上创建所以优化器状态天然跟着参数走不需要像旧方案那样手动去分片优化器状态。这算是我个人感觉最明显的一个变化心智负担降低了。接口上FSDP2把旧的wrap变成了一种显式的APIfully_shard。你不再需要定义什么“auto wrap策略”只要在模型构建时逐层调用它即可。比如某个Transformer模型有若干个Block你可以在模型初始化后对每个Block调用一次fully_shard(block)最后对整个模型再调用一次。PyTorch会自动把每个Block的参数分到各进程上。我用一个模拟项目X一个约13B的decoder-only模型在8卡机器上预训练实测过迁移成本原本用FSDP1写的训练脚本改成FSDP2大概只花了一个下午。改动点集中在模型初始化和优化器创建两处其余前向、反向、梯度裁剪的逻辑几乎没有动。下面是一个简洁的初始化示意import torch import torch.distributed as dist from torch.distributed.fsdp import fully_shard dist.init_process_group(nccl) # 先在meta设备上搭模型避免每个rank都真实分配一次参数 with torch.device(meta): model build_transformer(vocab_size50304, hidden_size5120, num_layers32) # 按业务上的模块边界做分片 for block in model.blocks: fully_shard(block) # 最后分片整个模型使其能在多卡间分布embedding、lm_head等剩余参数 fully_shard(model) # 从meta设备变成真实显存上的参数 model.to_empty(devicecuda)这里有一个关键点先构造meta设备上的模型再实际搬运是为了避免每个进程各自完整初始化一遍大模型省掉大量内存和时间。FSDP2在to_empty时会把每个参数只物化在自己所属的分片上所以搬运速度很快。2.2 分片粒度与通信调度的取舍FSDP2把“每个参数单独分片”作为核心设计这个选择和按层分片有什么区别最直接的影响是存储效率。按层或按整个模块分片意味着先把层内所有参数拼接成一个连续的buffer拼接过程会多一份临时拷贝另外许多层参数是非连续的强行拼成一段会留下pad空间。按参数分片则更加精确每个参数能算清楚自己在哪张卡上、以什么偏移量存放。对于不同shape的权重比如embedding、attention的q/k/v矩阵、MLM的头部它们的切分边界完全独立。实际显存开销更容易预测也避免了FSDP1里由于FlatParameter引起的“总显存没问题、单卡却偶尔OOM”的怪现象。通信调度上FSDP2遵循的是一个经典的两阶段模式前向阶段每个参数在计算前通过all-gather从分片卡收集到本卡上计算完后该参数立刻释放。后向阶段梯度计算依赖参数所以需要再次all-gather参数每张卡只得到自己分片对应的梯度slice然后通过reduce-scatter把各卡梯度归并到拥有对应分片的卡上。这样整个训练过程中的通信量大致是模型大小ψ的两倍前向一次all-gather约ψ后向又一次all-gather约ψ梯度reduce-scatter约ψ合计约3ψ其实精确来说all-gather和reduce-scatter各自传输的量分别是ψ所以FSDP单次迭代的通信量大约在2ψ到3ψ之间具体取决于实现细节和是否开启了梯度累积。相比DDP每次迭代通信2ψ的梯度交换FSDP并没有数量级上的增加却换来了“存储省好几倍”的收益这正是它被广泛采用的原因。不过有个容易被忽略的点FSDP2的通信是“每参数”级别的这意味着通信可以和计算重叠得更好。某个参数在前向结束后它的next参数就可以提前发起all-gather反向同理。在H100这类支持NVLink和IB高速互联的集群上重叠效果非常明显带宽利用率很高。2.3 HSDP当每张卡都觉得all-gather太贵时如果只是单纯分片全集群所有卡都要参与每份参数的all-gather。节点间的通信带宽通常比节点内的NVLink低好几倍一旦模型大到必须跨节点分片跨节点通信就会成为瓶颈。FSDP2支持一种叫HSDPHybrid Sharded Data Parallel的策略。思路很简单把并行维度分成两个——一个用于分片一个用于数据并行。例如在8卡机器上每台机器4卡你先在节点内做参数分片这4张卡分片保存同一份参数再在节点之间做普通的数据并行。参数的all-gather就只发生在节点内部4张卡之间跨节点只交换梯度聚合结果这样跨节点通信量大大降低而显存仍然只有原来的1/4左右。HSDP在实际训练中特别适合“一张节点内几张卡、节点间网络一般”的集群。我建议你在使用FSDP2时仔细看下fully_shard里的process_group参数将数据并行组和分片组分开设置。实测下来在同样集群上用HSDP替代纯FSDP后通信等待时间能降低30%以上而显存占用几乎不变。需要注意的是HSDP对梯度同步的语义有要求梯度先经过reduce-scatter到分片组内的对应卡再在数据并行组做all-reduce。这套流程PyTorch已经封装好但你要确保每张卡上的数据加载是按数据并行维度划分的这样每张卡算到的梯度才有意义。3. DeepSpeed ZeRO与FSDP2的异同3.1 ZeRO三阶段各在切什么DeepSpeed最重要的设计是ZeROZero Redundancy Optimizer。它把“参数量大”这件事按三个维度拆开每一阶段对应的显存节省完全不同ZeRO-1只对优化器状态做分片。每张卡只保存自己负责的那一部分Adam动量等状态但每张卡仍然保存完整模型参数和完整梯度。这适合那些“模型能放下、优化器状态放不下”的中间场景。ZeRO-2在ZeRO-1基础上再对梯度做分片。反向传播结束后各卡只保留自己分片对应的梯度部分其余梯度通过通信归并到对应卡上。这能使显存占用大幅下降。ZeRO-3再把模型参数本身也分片。前向和反向需要用到某个参数时通过broadcast或all-gather临时收集。这一步效果和FSDP是同一档的单卡不再需要完整存放整个模型。我第一次看到ZeRO-3时心里想的是“这不就是FSDP嘛”。确实从显存效果上看两者几乎等价。但从工程实践角度DeepSpeed有几个FSDP早期版本没有的差异化能力CPU Offload和NVMe OffloadZeRO-3可以把参数、优化器状态卸载到CPU内存甚至NVMe固态盘上这让单机训练几十B模型成为可能。缺点是整体训练吞吐会被磁盘读写拖累但作为“不得已而为之”的选择很实用。细粒度参数收集调度DeepSpeed支持按bucket的prefetch策略提前预取下一批参数降低等待时间。与HF生态的深度绑定大批预训练模型的加载、保存、断点续训在DeepSpeed里都有现成方案。3.2 通信量对比别被“零冗余”这个名号带偏“Zero Redundancy”指的是显存中的冗余为零不是通信为零。ZeRO-3在每轮前向、后向中同样需要把分片的参数收齐所以通信模式与FSDP基本一致前向all-gather一次、后向all-gather一次、梯度reduce-scatter一次。有人会问那是不是FSDP和ZeRO-3完全一样也不是。工程实现上的差异会造成实际性能差距通信算子是否和计算算子流水线式重叠。分片粒度DeepSpeed早期版本按“可能覆盖多个参数的buffer”分片FSDP2按“单个参数”分片前者在极端情况下会有更多的等待时间。梯度累积、offload、NVMe路径上的优化策略不同。如果强行排个性能结论在相同集群、相同模型、不开启offload的前提下FSDP2和最新版ZeRO-3的吞吐差距通常在5%以内很多时候取决于你有没有正确设置通信组和bucket大小。真正拉开差距的场景在于ZeRO在跨节点、多级分片、offload混合调度上更成熟而FSDP2在原生PyTorch生态里写起来更舒服。3.3 工程细节上的真实差异从使用者的感受来说FSDP2和DeepSpeed最大的差异在于“谁来管理生命周期”。DeepSpeed在使用时要维护一个deepspeed_config.json里面的字段包括zero_optimization的stage、offload_optimizer的device、通信bucket的大小等等。训练脚本需要把model丢给deepspeed.initialize来包装模型的前向、反向、step都要走DeepSpeed的封装。好处是配置化程度高坏处是一旦出错日志里的堆栈会套上厚厚一层封装。FSDP2则更像是“PyTorch原生的延展”它尽量不改变你的前向代码模型仍是一个普通nn.Module优化器也还是torch.optim里的那一个。它把“分片”“收齐”“释放”这些动作隐藏在模块钩子里。我个人的感受是如果团队里都是PyTorch流派的同学FSDP2上手会更快如果团队很习惯通过配置调参、并且需要offload能力那DeepSpeed更顺手。这一点没有绝对优劣只关乎组织的技术惯性。4. Megatron-LM的核心思路与互补4.1 张量并行把一个矩阵拆到多张卡上Megatron-LM走的是计算切分的路子。拿一个简单的线性层Y XW来说张量并行可以分为列并行和行并行。列并行时把权重矩阵W按列切成两份让两张卡各自拿一半输入X被广播到两张卡上各自算出一半的Y后再拼接起来。行并行时把W按行切两张卡各拿一部分行输入X也被切成两份分别送到两张卡最后把两边部分和的结果相加。这种切法带来的通信开销是每次矩阵运算前后都要做一次all-reduce或all-gather。在Transformer里attention层和MLP层各有一组这样的通信点所以层数越深通信次数越多。好处是每张卡只需存下模型一部分权重可以容纳极大的模型。Megatron的另一项优化是序列并行Sequence Parallel。普通张量并行中LayerNorm和Dropout仍然需要完整序列数据显存没有被切下去序列并行把序列维度也切开让这些逐点操作的显存进一步下降。这一套组合拳让它在超大模型上有很强的竞争力。4.2 流水线并行把网络切成多段Megatron还提供了流水线并行。简单来说模型的前若干层放在rank0中间若干层放在rank1依此类推。数据在多个stage之间传递时以micro-batch为单位形成1F1B一个前向接一个后向的调度模式让流水线尽量不断流。流水线并行的通信量不大因为只在相邻stage之间传activation和梯度但它对stage间带宽的延迟比较敏感。PP的调度本身有气泡时间bubble即流水线首尾造成的空闲所以单独用PP往往效率不高实践中几乎总是把TP、PP、DP组合使用。Megatron的优势在于它把这三者无缝集成在同一个框架里。你可以通过参数指定--tensor-model-parallel-size 4 \ --pipeline-model-parallel-size 2 \ --data-parallel-size 8再配合梯度累积形成一种“DP外、PP中、TP内”的层级结构。这种组合是数十B到数百B参数模型训练的常见配方。4.3 Megatron与FSDP2怎么配合很多人以为FSDP2和Megatron是竞争关系实际上在真正的工业级大规模训练中两者往往是组合使用的。经典的组合方式有TP并行在节点内因为张量并行需要极高的带宽放在单机内部通过NVLink通信最合适。PP并行在节点间把流水线的不同stage放到不同机器上通信量不高对网络带宽要求相对低。FSDP/ZeRO负责数据并行维度多个“TPPP复制组”之间再做数据并行这时每份模型副本的参数可以进一步用FSDP分片从而让“每组一张卡只放一种参数”变成“每组多张卡共同分摊一份模型”。举个例子一个70B模型8卡机器共8台机器64卡。你可以在每台机器内开TP4将模型切成两段做PP2于是每台机器形成一个TP4、PP1的模型副本4个副本之间用HSDP/FSDP做数据并行。这样每一份分片后的模型参数在这4个副本之间又切了一道单卡实际保存的参数远小于原始模型。我用模拟项目X做过类似实验结论是只靠FSDP264卡训练70B模型可行但跨节点通信压力不小叠加TP后在单卡上的显存占用更低训练吞吐能再提升10%以上。反过来只在Megatron体系里用TPPP而不加FSDP/data parallel分片代码复杂度会上升很多尤其是处理动态batch、梯度裁剪、混合精度时都要小心兼容。所以我的建议很明确模型在30B以下、团队追求简单优先FSDP2模型在30B以上、机器规模超过32卡考虑Megatron框架或“TPPPFSDP”组合方案DeepSpeed则作为一套成熟的parameter sharding offload替代方案。5. 框架选型预算、规模与团队能力5.1 一个可复用的选型决策流程面对具体项目时我不会一上来就讨论哪个框架更好。更合理的顺序是先算一笔账模型显存需求是多少把参数、梯度、优化器状态分开计算明确必须在多少张卡上分片。节点内、节点间的网络拓扑如何有没有NVLink、节点间是100Gb还是200Gb的IB如果节点间带宽一般务必给HSDP或DeepSpeed的多节点分片留好通信组。训练代码的侵入性要求高不高如果现有训练脚本是纯PyTorch写的FSDP2迁移成本最低如果项目本来就用了HF的Trainer那么HF对DeepSpeed的集成会更友好一些。是否需要offload显存实在不够、或者只有单机多卡时DeepSpeed的ZeRO-Offload比FSDP2自带的东西更成熟。是否要大规模扩展如果团队计划在几百张卡级别上长期训练超大规模模型越早切到“TPPP数据并行/FSDP”的体系越划算。把这些问题过一遍答案基本就出来了。5.2 三套方案对比速查表下面这个表是我在几个项目里总结出来的可以当作参考对比维度FSDP2DeepSpeed ZeROMegatron-LM核心思路参数/梯度/优化器状态分片存储视角去冗余参数/梯度/优化器状态分片并支持CPU/NVMe offload张量并行、序列并行、流水线并行计算视角切分与PyTorch兼容性原生、几乎无侵入有封装需经DeepSpeed引擎初始化需改写模型为Megatron的分层结构易用性高适合大多数团队中高配置调优成本中低需要并行思维和更细致的调试单卡显存优化好好offload能力更强好但主要面向TP/PP场景通信开销中等与ZeRO-3相当中等可精细调优高TP通信频繁但对大模型收益巨大多节点训练支持HSDP跨节点友好支持多节点、多级分片支持但节点间建议跑PP/DP典型适用场景中小团队、快速迭代、模型几十B需要offload、单机/小集群、HF生态超大模型、几十到几百B、规模化生产训练5.3 混合使用案例70B模型在64卡集群上的配置这里给一个可复现的思路。假设某实验室的机器是8台8卡设备要训练一个70B decoder模型目标吞吐约2k tokens/s/GPU以内显存完全放不下网络为NVLink200Gb RoCE。节点内TP4把attention和MLP的矩阵切两半。这样每张卡只负责半个矩阵激活显存大幅下降。节点内PP2把前一半层放在前4卡后一半层放在后4卡。节点间DP4也就是4个“TP4×PP2”的副本每组负责不同的数据。这一层如果再加FSDP2则可以不再重复保存完整模型副本。优化器状态加载完成后用CPU offload或FSDP分片再做一次分片。这套方案在训练时每张卡存下的参数大约是70B / 64 ≈ 1.1BFP16下约2.2GB优化器状态再分片后也只有几GB完全在卡容量范围内。通信上TP的all-reduce在节点内部用NVLink完成跨节点的FSDP梯度聚合在数据并行组内完成负载更均衡。我用这套思路在一个类似规模的实验里把显存峰值压到了70%以下训练稳定性也可以接受。但提醒一句TP4时每个Transformer层的计算量变小通信占比会上升所以并不总是越大越好。TP2或TP4是一个常见甜点具体还要看卡型和模型size。6. FSDP2实战踩坑与问题排查6.1 从裸训练迁移到FSDP2的第一天如果你的项目目前还是“单卡或DDP”的训练脚本想把模型换成FSDP2训练我建议按下面的顺序改造第一步把模型放置改成meta device。在build_model时不使用真实的随机初始化而是先通过with torch.device(meta)把结构建出来。这能让你在代码里“看到”模型的所有参数但不占显存。第二步逐模块调用fully_shard。不要只对最外层模型调用一次那样会让所有参数挤在一个大分片组里单卡仍然要完整收集模型才能计算。合理的做法是让每个Transformer Block成为分片单元让embedding层、lm_head单独分片。第三步再创建优化器。因为参数已经被物理分片了优化器只会看到本卡拥有的分片参数所以优化器状态也天然是分片的。不要使用torch.compile直接包一个已经分片过的模型虽然新版本兼容性变好了但调试时先关闭compile能减少变量。第四步设置clip_grad_norm_时要小心。由于梯度可能是分布式的直接调用普通的梯度裁剪会得到错误结果。在FSDP2里建议使用torch.nn.utils.clip_grad_norm_时传入所有模型参数或者使用FSDP专门的处理方式否则裁剪的量会被低估。第五步checkpoint的保存加载。FSDP的参数分片后每张卡只保存自己那部分。如果你想得到完整权重需要在gathered状态下保存如果想省时间可以每个分片各存一份加载时由框架自动组合。两种模式我都有用过纯速度上分片保存快得多但生成权重文件需要额外做一次合并操作。6.2 常见问题速查表现象大概率原因解决办法单卡OOM但总显存看起来够忘记开启HSDP、所有参数都在同一分片组给fully_shard传不同的process_group按节点内和节点间分层训练速度比DDP还慢通信没有与计算重叠、bucket尺寸不对检查是否开启了FSDP的use_orig_params逻辑或尝试把分片模块粒度调大梯度不准或模型不收敛梯度裁剪出错、混合精度下master weight未分片使用FSDP上下文中的clip_grad_norm_确认no_sync与梯度累积的配合跨节点训练时频繁timeout节点间all-gather通信等待且没有HSDP用HSDP降低跨节点通信频率调整NCCL_TIMEOUT等环境变量但不要盲目加大保存模型后load回来自变量乱掉checkpoint下的模型状态key命名被FSDP的封装改了用state_dictAPI重新保存一次不要直接保存原始module的dict6.3 框架切换时容易踩的坑从DeepSpeed切到FSDP2时最容易出问题的是学习率缩放。DeepSpeed在ZeRO-3下对参数做了额外处理同样的显存占用可能会改变batch size对梯度的统计FSDP2下的梯度是真正的全局平均梯度切换后容易出现loss波动。建议切换后先用小学习率跑几十个step对比loss和梯度范数。从Megatron切到FSDP2时最大的痛点在于模型结构的差异。Megatron为了TP做了权重切分和fusion模型定义本身就带并行语义FSDP2希望模型保持原始PyTorch结构。如果你手里有一份Megatron化的模型权重直接load到FSDP2会很痛苦。要么从原始权重预训练保存版本开始要么写一份权重重组脚本把切分过的权重拼接回来。7. 最后谈一点选型体会我自己先后经历过几轮框架迁移最早是单卡实验后来上DDP再后来因为模型到了几十B级别不得不逐层学习FSDP的动态分片机制又因为需要offload尝试了DeepSpeed最后为了验证超大规模场景还写了不少Megatron相关的实验脚本。走到今天我的体会是没有一个框架能同时做到“最省显存、最快速度、最简单上手”三个目标。FSDP2的定位很清晰它是PyTorch生态里最顺手、最通用的参数分片方案适合大多数“先跑起来再优化”的场景DeepSpeed在offload和细粒度通信调度上经验老到适合需要压榨单机显存的团队Megatron则更接近“为几百卡以上生产环境设计”的重型框架学习曲线陡但一旦跑顺稳定性和扩展性都很可靠。如果你现在正纠结选哪个不如先回答一个问题这个项目是三天内要看到训练曲线还是三个月内要稳定产出大规模模型权重。前者选FSDP2后者建议认真考虑Megatron或者“TPPPFSDP”的组合路线。另外不管选哪套方案我都建议在做大规模训练之前先用一个小模型在小规模集群上把通信拓扑和checkpoint恢复的逻辑测试几轮这几小时的时间投入往往能帮你避开后面几天的排障苦战。我在实际使用中踩过最深的坑是看文档时觉得FSDP2很简单就直接把70B模型的脚本从DDP切成FSDP2结果在一周里反复遇到奇怪的通信死锁和显存碎片。后来冷静下来先拿7B模型在4卡上把前后向、梯度裁剪、混合精度、断点续训全部调通再放大到全量模型整个过程反而顺畅得多。如果你正打算大规模切FSDP2不妨也先走一遍这条“小模型验证、大模型跑量”的路。