Data Parallel

设第 $r$ 个 rank 有 $n_r$ 个有效 token，局部平均梯度为 $g_r$ 。全局 token 平均梯度应按 $n_r$ 加权。采用默认按 rank 平均的 DDP 时，可以把局部平均 loss 乘以 $p n_r/\sum_j n_j$ ，其中 $p$ 是参与同步的 rank 数。

梯度累积时，分母还应覆盖整次更新中的所有有效 token。计算这个计数时，padding 与被 mask 的位置不能算进有效 token。

g=\frac{\sum_r n_rg_r}{\sum_r n_r}

DDP API 不等 batch 讨论 #

ZeRO 的三个阶段分别省掉了什么？

ZeRO-1 分片优化器状态。ZeRO-2 进一步分片梯度。ZeRO-3 再分片参数，计算某部分模型前按需收集。各阶段都在减少数据并行副本之间重复保存的模型状态。

FSDP 的完整分片模式与 ZeRO-3 有相近的状态分片目标。实际峰值还包括激活、临时收集的参数、预取和通信缓冲区，不能把整张卡的占用直接除以并行度。

ZeRO FSDP2 教程 #

FSDP 和 TP 都分片权重，为什么还要区分？

常见 FSDP 会在计算前收集当前模块需要的完整参数，各数据 rank 用它处理自己的输入。TP 则在计算期间继续持有权重分片，由多张卡共同完成同一层的计算。

这使它们的通信对象和时机不同：FSDP 主要围绕参数与梯度的生命周期，TP 主要围绕层内计算与激活布局。两者可以在不同 DeviceMesh 维度组合。

FSDP2 教程 TorchTitan FSDP #

DDP 与完整分片训练，通信量怎样比较？

设同步组有 $p$ 个 rank，完整模型参数和梯度的有效大小都为 $M$ 字节。先用 Ring 算法，只统计每个 rank 的发送量。DDP 一次梯度 AllReduce 约为 $2(p-1)M/p$ 。

若完整分片训练在前向后释放完整参数，并在反向前重新收集，一轮典型过程包含两次参数 AllGather 和一次梯度 ReduceScatter，约为 $3(p-1)M/p$ 。

参数与梯度精度不同、保留完整参数、重计算或混合分片都会改变这个计数。字节数也不直接决定步时：需要继续看哪些通信能与计算重叠。

\begin{aligned}V_{\mathrm{DDP}}&\approx2\frac{p-1}{p}M\\[4pt]V_{\mathrm{FSDP}}&\approx3\frac{p-1}{p}M\end{aligned}

NCCL Tests FSDP2 教程 fully_shard API #

梯度累积时，为什么常常先 no_sync，最后再同步？

在同一次参数更新内，梯度可以先在本地累积，再统一同步，减少重复的梯度归约。DDP 的 no_sync 需要包住对应的前向和反向，最后一次正常的前向与反向触发同步。

累积的 loss 仍需按这次更新的有效 batch 归一化。FSDP 的梯度同步与参数收集是不同控制项，关闭梯度同步并不自动消除参数 AllGather，也可能增加梯度的显存占用。

DDP API fully_shard API #

为什么 HSDP 要在一组卡内分片，再在组间复制？

在较小的组内分片，可以把频繁的参数收集限制在较快的互连内。不同组各自处理数据，再同步对应的梯度分片。

代价是模型状态会在多个组之间重复保存。它用更多显存换取不同的通信布局，是否更快取决于模型、batch 和节点内外带宽。

HSDP 图解 fully_shard API #

推理里的 DP，也要像训练一样同步梯度吗？

推理不计算用于参数更新的梯度。普通稠密模型的多个服务副本可以各自处理不同请求，主要目的是扩展吞吐和并发容量。

训练中的 DP 则需要把不同数据产生的梯度合起来。两者都把输入分给多份模型，但不能把训练的梯度同步成本直接算到推理副本上。

PyTorch Distributed #

FSDP 预取更积极，怎样判断省下的等待是否值得额外显存？

待解

同时观察完整参数的存活区间、通信缓冲区、激活峰值和未被计算覆盖的 AllGather。最合适的分组大小会怎样随模型和 batch 改变？

FSDP2 教程 TorchTitan FSDP FSDP2 验证 PR #