FSDP与ZeRO：分布式训练对比

可视化演示

对比FSDP与ZeRO的分片策略、显存优化阶段、通信模式与框架选型。

分布式训练背景：为什么一个人搬不动

老陈接到达康图书馆整体搬迁的任务，书库里几百万册书，他请来八位搬家师傅。问题很快暴露：书架太大太重，任何一位师傅单独都搬不下。大模型训练也是如此——当模型参数涨到百亿级别，单块 GPU 的显存根本装不下完整的权重、梯度和优化器状态，这就是“显存瓶颈”。

于是有了“并行”的思路。最朴素的是数据并行（Data Parallelism）：每位师傅都备一套一模一样的书架，各自扛到不同楼层，开工前用对讲机对齐进度。对应到训练中，每张 GPU 都保存完整模型副本，处理不同 batch 的数据，最后把梯度汇总同步。它简单好用，但显存代价是“每人复制一整套家具”。另一种是模型并行（Model Parallelism）：把书架本身拆成几段，每人扛一段，缺点是楼层间搬运必须层层等待，串行通信多。混合并行则是两者叠加，用数据并行铺开、模型并行切深。

数据并行下，每张 GPU 需要保存四份关键状态：一份模型参数、一份反向传播算出的梯度，以及 Adam 优化器里保存的两份动量状态。用 P 表示一份参数所需的显存，单卡需求约为 4P，公式化的表达是：

$$\displaystyle M_{\text{DDP}} = P_{\text{参数}} + P_{\text{梯度}} + 2P_{\text{优化器}} = 4P$$

1. 参数量 P：$P$ 表示一份模型权重占用的显存空间，7B 模型仅参数就需要约 28GB。

2. 参数副本：$P_{\text{参数}}=P$ 当前权重本身，每张卡都持有一份。

3. 梯度副本：$P_{\text{梯度}}=P$ 反向传播给每个参数算出的修正值，和参数一样大。

4. 优化器状态：$2P_{\text{优化器}}=2P$ Adam 中动量 m 和二阶矩 v 各占一份，是训练优化器必需的“保养手册”。

数据并行虽然简单，却像一个“人人各自背全套”的搬家队——八位师傅每人背同样重的柜子，显存浪费呈线性增长。模型稍大一点就超出单卡极限。要突破这个瓶颈，必须换一种思路：与其人人扛全套，不如把柜子拆开，八位师傅各扛一条腿。这正是 ZeRO 的思想萌芽。

ZeRO分片理念：零冗余的搬家队

老陈重新设计搬法：整柜子一个人扛不动，那就拆成八份，大家只扛一份。每个师傅的负重降到原来的 1/8。ZeRO（Zero Redundancy Optimizer，零冗余优化器）的名字就是这套思路的宣言——把数据并行中“每张卡都存一遍”的重复冗余全部去掉。

ZeRO 将冗余消除分三个阶段推进。Stage 1 先分片优化器状态，也就是把 Adam 的动量“保养手册”拆开，每人只保存自己负责页面的那几页；Stage 2 再把梯度也分片，每个人只保留更新自己那部分参数所需的修正值；到了 Stage 3，连模型参数本身都拆开存放，前向传播要用时再通过通信临时拼装。三阶段合在一起，就是完整的 ZeRO-3，也是 DeepSpeed 引擎默认支持的分片模式。

在 ZeRO-3 中，参数、梯度、优化器状态在 N 张 GPU 间平均分配，每张卡只保留总量的 1/N：

$$\displaystyle M_{\text{ZeRO-3}} = \frac{P + P + 2P}{N} = \frac{4P}{N}$$

1. 分子四项：$P + P + 2P$ 把参数、梯度、优化器状态全部拆开，它们不再每卡各存一份。

2. 分母 N：$N$ 是参与训练的 GPU 总数，所有状态被平均切分到每张卡上。

3. 合并化简：$\frac{4P}{N}$ 显存随卡数线性下降，这是 ZeRO-3 最诱人的地方。

老陈把这套“拆扛方案”在团队里推广：没有冗余，每个师傅负重轻了 8 倍。但正如拆家具需要把零件交给正确的人，训练也得先把拆开的参数在每一层重新拼合起来。这份“重新拼合”的代价，就是通信成本。

FSDP核心机制：PyTorch 里的“同款制度”

老陈的方案很优秀，但 PyTorch 团队也推出了自己的版本：FSDP（Fully Sharded Data Parallel，全分片数据并行）。它和 ZeRO-3 几乎同源，把参数、梯度、优化器状态全部摊到各卡上，只是 API 和调度方式不同。就像两家公司都采用“分拆扛”制度，但一个用纸笔记录分配，一个用 App 自动排班。

FSDP 的分片粒度更细——它按“层”切分，而不是把整个模型当成一整块。模型有多少层，就安排多少轮“借还”：训练到第 i 层时，持有该层参数的 GPU 通过 all-gather 把参数广播给全体；算完这一层，又把参数碎片归还。反向传播时，同样的收集过程还要再来一次，确保梯度能被正确计算。

FSDP 的核心通信量可以拆成三项：前向为计算而收集一次参数（约 2P），反向为算梯度再次收集（约 2P），最后用 reduce-scatter 把梯度汇总回各自分片（约 2P）。三项相加大约等于 DDP 的三倍：

$$\displaystyle C_{\text{FSDP}} \approx 2P_{\text{前向}} + 2P_{\text{反向}} + 2P_{\text{梯度}} = 6P \approx 3 \times C_{\text{DDP}}$$

1. 前向收集：$2P_{\text{前向}}$ 前向传播时需要完整参数，被分片的参数要通过 all-gather 收集一遍。

2. 反向再收集：$2P_{\text{反向}}$ 反向传播又要用到参数，没有缓存就得再收集一次。

3. 梯度归约：$2P_{\text{梯度}}$ 反向产生的梯度经过 reduce-scatter，每个节点只拿到自己负责的那份。

4. 对比基准：$3 \times C_{\text{DDP}}$ DDP 一次 all-reduce 约 2P，FSDP 的通信总量约为其三倍。

站在老陈的角度，FSDP 是一套自动化的“分拆扛制度”：不需要手动告诉每一层怎么分，配置好 wrap 策略，PyTorch 会自动完成拆解和交接。它让 PyTorch 用户花更少精力就获得接近 ZeRO-3 的显存收益。

通信模式对比：对讲机里的两句话

“分拆扛”解决了负重问题，也带来新问题——频繁的对讲机沟通。搬家时从一个师傅手里拿书要喊，还回去要喊，谁的意见综合完还要对着所有人广播一次。训练中的通信原语很像这几句“喊话”，其中最常用的是 all-gather 与 reduce-scatter。

all-gather 是“众筹式收集”：每个节点把自己手头的参数碎片发给所有其他人，最终每个节点都获得完整参数。reduce-scatter 是“汇总式分配”：所有节点先各自算出梯度碎片，经过两两交换核对求和，最终每个节点只保留自己负责的那一份。一个负责“凑整”，一个负责“分摊”。

通信次数和消息大小共同决定开销。工程上常用一个简单的延迟模型来估算单次通信时间：

$$\displaystyle T_{\text{comm}} = \alpha + \beta \times S$$

1. 启动延迟 α：$\alpha$ 发起一条消息的固定开销，类似打电话时的拨号等待时间。

2. 带宽系数 β：$\beta$ 每传输 1 字节数据需要的时间，由网卡和拓扑决定。

3. 消息大小 S：$S$ 单次通信传输的字节数，例如 P 个参数对应的总字节量。

4. 结合含义：$T_{\text{comm}}= \alpha + \beta \times S$ 总通信时间 = 固定开销 + 传输开销，碎消息越多越吃亏。

老陈发现，对讲机喊话虽然花时间，但如果一边搬上一个人，一边提前联系下一个，就能把沟通时间藏进搬运动作里。FSDP 的“计算通信重叠”正是这个思路：把通信切成小块，插进层与层的算子间隙中，让 GPU 计算和网卡传输并行发生，从而掩盖大部分通信延迟。

JAX实践要点：用“分片说明书”自己组装

老陈的公司新购入了一套万能对讲系统——JAX。它不内置现成的“搬家方案”，但提供一整套“分片说明书”API：你想让哪个师傅负责哪些书、在哪个楼层、用哪条通道，都可以精确声明。JAX 信奉显式优于隐式，因此把分片控制权完全交给你。

JAX 里没有名为 FSDP 的类，但有三个核心概念。Mesh（设备网格）把 GPU 排成多维矩阵，给它一个逻辑轴名字；PartitionSpec（分片规格）声明张量的每个维度落到 Mesh 的哪个轴；最后用 jax.jit（旧称 pjit）绑定输入/输出分片规则并编译。三者组合就能复现 ZeRO-3 式的效果，任何框架层无法提供的定制空间都可以在这里实现。

以把张量的第一维按行切到 4 张卡为例：Mesh 定义一个叫 "fsdp" 的轴，PartitionSpec('fsdp', None) 表示第一维分片、第二维不切。用数学语言描述，第 i 行的数据会归于编号为 $i \bmod N$ 的设备：

$$\displaystyle X[i,\; :] \to \text{GPU}_{(i \bmod N)}$$

1. 张量切片：$X[i,:]$ 取张量的第 i 行，整行数据是分片的最小单元。

2. 设备编号：$\text{GPU}_{(i \bmod N)}$ 用 mod 计算第 i 行归哪台 GPU，保证轮转分布。

3. 分片语义：$i \bmod N$ 把连续行切给不同设备，等价于 PartitionSpec 的第一维分片。

老陈觉得这套系统的关键是“显式”：没有框架替他做决定，但一旦配好 Mesh 和 PartitionSpec，训练便按声明自动分布式执行。对有定制需求的团队，这种自由度比“拿来即用”更宝贵，尤其适合以 JAX 为底座的模型研究团队。

框架选型指南：到底选哪家的“分拆扛”

方案大同小异，选哪家却要看场景。老陈给几家客户做了对比：PyTorch 用户多、要快速上线，选 FSDP；DeepSpeed 用户看重极致显存优化和 CPU/GPU 卸载，选 ZeRO；JAX 生态研究者则更习惯自己装配。这三条路没有绝对优劣，只有匹配不匹配。

先看 PyTorch FSDP——它是 PyTorch 原生的 API，零额外依赖、配置简单，适合中小团队快速启用，也易于与 torch.compile、DDP 无缝衔接。再看 DeepSpeed ZeRO——它是微软开源的引擎，支持更细粒度的 stage 控制、Offload 到 CPU 甚至 NVMe，以及混合精度调度，但要引入额外依赖和配置。JAX 则是另一种风格：用 jax.sharding 手动编排分片规则，灵活度最高，但也要求使用者对通信模式有深层理解。

当模型规模确定时，可以用一个经验公式估算最少需要的 GPU 数量。设单卡可用显存为 $M_{\text{GPU}}$：

$$\displaystyle N_{\text{min}} = \left\lceil \frac{4P}{M_{\text{GPU}}} \right\rceil$$

1. 分子：$4P$ 包含参数、梯度、优化器状态，是 DDP 模式下的总显存需求。

2. 分母：$M_{\text{GPU}}$ 单张 GPU 可用于训练状态的显存容量。

3. 向上取整：$\lceil \cdot \rceil$ 保证即使非整数也按整卡数配置，卡数不足则跑不起来。

老陈最后给客户的建议是：40GB 显存、8 卡以内的中小规模，优先 FSDP；千亿级模型或需要细化卸载策略，DeepSpeed 更合适；如果团队本身用 JAX 生态（如 T5X、Flax），直接用 jax.sharding 实现分片。分片思想都一样，选型只取决于你站在哪个生态里。

AI 应用场景

Python 演示

def estimate_memory_per_gpu(param_count, num_gpus, stage):    """    估算不同分片策略下每张 GPU 的最低显存需求。    单位: GB, 只考虑参数/梯度/优化器状态, 不含激活值。    """    param_bytes = param_count * 4          # float32 参数, 每个占 4 字节    grad_bytes = param_count * 4           # 梯度, 形状同参数    opt_bytes = param_count * 4 * 2        # Adam 的 m 和 v, 各占一份     if stage == 0:              # 数据并行 DDP        per_gpu = param_bytes + grad_bytes + opt_bytes    elif stage == 1:            # ZeRO-1: 仅优化器状态分片        per_gpu = param_bytes + grad_bytes + opt_bytes / num_gpus    elif stage == 2:            # ZeRO-2: 优化器状态 + 梯度分片        per_gpu = param_bytes + grad_bytes / num_gpus + opt_bytes / num_gpus    elif stage == 3:            # ZeRO-3 / FSDP: 全部分片        per_gpu = (param_bytes + grad_bytes + opt_bytes) / num_gpus    else:        raise ValueError("stage 仅支持 0-3")    return per_gpu / (1024 ** 3)  def train_demo():    param_count = 7_000_000_000  # 7B 模型    gpus = 8    stage_names = {        0: "DDP (数据并行)",        1: "ZeRO-1",        2: "ZeRO-2",        3: "ZeRO-3 / FSDP",    }    print("7B 模型在不同的分片策略下, 每张 GPU 的最低显存需求:")    for stage, name in stage_names.items():        gb = estimate_memory_per_gpu(param_count, gpus, stage)        print(f"{name:16s} 约 {gb:6.1f} GB")  if __name__ == "__main__":    train_demo()

1. param_bytes = param_count * 4一个 float32 参数占 4 字节，参数总量决定后续所有状态大小。
2. opt_bytes = param_count * 4 * 2Adam 优化器要维护 m 和 v 两组状态，各占一份参数空间。
3. per_gpu = param_bytes + grad_bytes + opt_bytesDDP 模式下三份状态每卡完整复制，这就是 4P 的本质。
4. per_gpu = param_bytes + grad_bytes + opt_bytes / num_gpusZeRO-1 只把最占空间的优化器状态摊到各卡。
5. per_gpu = param_bytes + grad_bytes / num_gpus + opt_bytes / num_gpusZeRO-2 进一步分片梯度。
6. per_gpu = (param_bytes + grad_bytes + opt_bytes) / num_gpusZeRO-3/FSDP 三份状态全部均摊，显存随卡数线性下降。

常见误区

面试题