FSDP与ZeRO3谁更强？ms-swift框架下大规模训练实测结果-开发者社区

FSDP与ZeRO3谁更强？ms-swift框架下大规模训练实测结果

在大模型时代，显存墙问题日益突出——一个70B参数的LLM全精度训练动辄需要数百GB显存，远超单卡容量。面对这一挑战，分布式训练不再是“可选项”，而是“必选项”。而如何在有限资源下高效训练百亿甚至千亿级模型，FSDP（Fully Sharded Data Parallel）和ZeRO-3（Zero Redundancy Optimizer Stage 3）成为最常被提及的两种解决方案。

两者目标一致：打破数据并行中冗余状态复制带来的显存爆炸。但路径迥异——一个是PyTorch原生轻量派，另一个是DeepSpeed驱动的重型引擎。究竟谁更适合你的场景？本文基于ms-swift框架的实际支持能力，结合工程实践视角，深入剖析二者的技术差异、性能边界与选型逻辑。

显存优化的本质：从 $ O(3P) $ 到 $ O(P/N) $

传统数据并行（如DDP）每个GPU都保存完整的模型副本、梯度和优化器状态。以Adam优化器为例，每张卡需存储：

模型参数：$ P $
梯度：$ P $
一阶动量（momentum）：$ P $
二阶动量（variance）：$ P $

总显存消耗为 $ 4P $，即 $ O(P) $ 级别。当使用FP16混合精度时可压缩至 $ 3P $，但仍随模型规模线性增长。

FSDP与ZeRO-3的核心突破在于引入分片机制（sharding），将这些状态按设备数 $ N $ 分布式存储。理想情况下，每卡仅保留 $ 1/N $ 的参数、梯度和优化器状态，理论显存占用降至接近 $ O(P/N) $，实现近乎线性的显存扩展。

虽然数学表达相似，但两者的实现方式、系统依赖和工程代价截然不同。

FSDP：PyTorch原生的“轻骑兵”

FSDP由Meta提出，并自PyTorch 1.12起作为官方模块集成，定位清晰：为大多数用户提供一种无需额外依赖即可获得显著显存收益的方案。

它的工作模式可以用一句话概括：前向/反向传播时动态聚合所需参数，计算完成后立即释放非本地分片。

具体流程如下：

初始化阶段：模型被划分为多个子模块（例如Transformer层），每个模块独立包装为FSDP单元。
前向传播：
- 当前设备加载本地参数分片；
- 若需跨设备参数，则通过all-gather操作临时拉取完整权重；
- 完成计算后释放该部分内存。
反向传播：
- 计算局部梯度；
- 使用reduce-scatter将梯度归约到对应设备；
- 各设备仅更新自身持有的参数分片及其对应的优化器状态。
参数同步：下一轮前向传播前重新聚合。

这种“用通信换显存”的设计，在中小规模集群中表现稳健。尤其适合ms-swift倡导的“一键启动”理念——开发者无需搭建复杂的运行时环境，只需修改几行代码即可启用。

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP import torch.nn as nn model: nn.Module = get_llm_model() fsdp_model = FSDP( model, mixed_precision=torch.distributed.fsdp.MixedPrecision( param_dtype=torch.float16, reduce_dtype=torch.float16, buffer_dtype=torch.float16 ), cpu_offload=CPUOffload(offload_params=True), # 可选卸载至CPU device_id=torch.cuda.current_device(), )

这段代码展示了FSDP的典型用法。通过设置mixed_precision启用FP16/BF16加速，配合cpu_offload进一步降低GPU压力。值得注意的是，FSDP支持逐层控制是否启用分片（viaauto_wrap_policy），便于对注意力层等高显存模块做精细化管理。

不过，这种灵活性也带来了潜在瓶颈：频繁的all-gather操作对通信带宽极为敏感。在千兆网络或多跳交换环境中，延迟可能成为制约吞吐的关键因素。

此外，FSDP目前不支持与Tensor Parallel或Pipeline Parallel深度协同，难以直接用于万亿参数级别的超大规模训练任务。

ZeRO-3：工业级训练的“重装部队”

如果说FSDP是一把灵活的手术刀，那么ZeRO-3更像一台重型挖掘机——专为极限场景打造。

作为DeepSpeed框架中的第三阶段零冗余技术，ZeRO-3实现了三重分片：

参数分片（Stage 3新增）
梯度分片（继承自Stage 2）
优化器状态分片（继承自Stage 1）

这意味着在整个训练过程中，没有任何一张GPU持有完整的模型副本。所有状态均按 $ 1/N $ 分布，真正做到了“去中心化”。

其工作流程如下：

参数被均匀切分并分布于各GPU；
前向传播时，若某层访问非本地参数，则触发broadcast或all-gather获取；
计算完成后立即释放，避免长期驻留；
反向传播中，梯度经reduce-scatter归约，仅更新本地对应的参数分片；
支持激活值检查点（activation checkpointing）、CPU/NVMe卸载等高级特性。

正是这套机制，使得ZeRO-3能够支撑如Megatron-Turing NLG（530B）这样的庞然大物。配合DeepSpeed的pipeline parallelism和tensor parallelism，可在数千卡集群上实现近线性扩展。

更重要的是，ZeRO-3提供了强大的offload能力。以下配置可将优化器状态和参数卸载至CPU甚至NVMe：

{ "train_batch_size": 8, "gradient_accumulation_steps": 4, "fp16": { "enabled": true }, "zero_optimization": { "stage": 3, "offload_optimizer": { "device": "cpu" }, "offload_param": { "device": "cpu" }, "overlap_comm": true, "contiguous_gradients": true }, "activation_checkpointing": { "partition_activations": true, "cpu_checkpointing": true } }

该配置启用了通信重叠（overlap_comm）、连续梯度缓冲等优化，最大限度掩盖I/O开销。在ms-swift中，用户只需指定此JSON文件，即可通过Deepspeed插件自动加载，大幅降低了使用门槛。

然而，这一切的前提是：你得有一套高性能IB/RoCE网络、充足的CPU内存和成熟的运维体系。否则，复杂的配置调试过程足以劝退多数小团队。

实战对比：不只是“谁更快”

在ms-swift框架的设计中，FSDP与ZeRO-3并非互斥选项，而是构成了一套覆盖全场景的分布式训练路由系统：

+-----------------------------+ | ms-swift CLI | | (yichuidingyin.sh 脚本入口) | +------------+---------------+ | +-------v--------+ +------------------+ | 用户配置选择 | --> | 分布式训练引擎路由 | +----------------+ +------------------+ / \ / \ +----------+--+ +-------+---------+ | FSDP Backend| | ZeRO-3 + DeepSpeed| +-------------+ +-------------------+ | | PyTorch Native DeepSpeed Runtime (Built-in) (External Lib)

根据用户指定的parallelism_type，系统自动切换底层执行引擎，屏蔽了API差异。但这并不意味着可以随意选择——真正的决策应基于实际需求权衡。

工作流程差异一览

阶段	FSDP	ZeRO-3
初始化	直接包装模型	加载DeepSpeed配置并初始化引擎
显存分配	参数/梯度/状态本地分片	全局分片 + CPU/NVMe offload支持
前向传播	All-gather拉取完整参数	动态广播或聚合，支持分区激活
反向传播	Reduce-scatter梯度	类似，但支持更细粒度通信调度
检查点保存	保存分片，合并需gather	支持`stage3_gather_16bit_weights_on_model_save`
推理兼容性	导出后可直接部署	必须先gather完整权重才能用于vLLM/LMDeploy