利用联邦学习加速大规模分片数据并行

论文信息

标题: Accelerating Sharded Data Parallelism at Scale with Federated Learning

作者: Gianluca Mittone, Marco Aldinucci

发布日期: 2026-09-17

arXiv ID: 2609.20359v1

PDF 链接: 下载 PDF

3 分钟速览

• 研究问题：大规模分片数据并行（FSDP/HSDP）在多级互连 HPC 系统上的通信开销随 GPU 数量急剧上升，标准 FSDP 难以扩展到数百至数千个 GPU。
• 核心方法：将联邦学习中的 FedAvg 式聚合作为更高层并行原语，与 FSDP/HSDP 交错，形成 FL+FSDP 和 FL+HSDP，把训练进程划分为松耦合的联邦组，降低跨组通信频率。
• 关键结果：在 512 个 A100 GPU 上预训练 Llama3.1 8B 时，FL+FSDP 的数据处理速度比 FSDP 快 8.04 倍；FL+HSDP 的评估困惑度比 HSDP 低 4.48 倍（表 3）。
• 主要局限：FL+HSDP 的吞吐优势在该实验设置下未兑现，作者推测原因是带宽差异不足与 SLURM 非确定性节点分配；FL 聚合步骤引入较高的执行时间方差。
• 适合读者：从事大规模模型训练、HPC 分布式系统、联邦学习与数据并行优化的研究者和基础设施工程师。

论文背景和研究动机

现代 AI 训练越来越依赖大规模 HPC 集群。基础模型的参数量和训练数据量持续增长，训练任务需要数千个 GPU 协同工作。数据并行是最常用的分布式训练策略，它让多份模型副本在不同数据批次上并行训练，并通过每批次的梯度同步保持副本一致。为了训练单卡放不下的大型模型，分片数据并行进一步将模型参数切分到多个 GPU 上，其中 Full Sharded Data Parallelism（FSDP）沿一维 sharding group 切分参数，Hybrid Sharded Data Parallelism（HSDP）则在二维 $R \times S$ 网格上同时使用分片和复制。

问题在于，现代 HPC 和数据中心的互连性能高度非均匀：节点内 GPU 间可以使用 NVLink 等高速链路，带宽可达约 1.8 TB/s；节点间通常依赖 InfiniBand 等慢一到两个数量级的网络，例如每链路约 200 GB/s；机架之间还会出现额外网络噪声，有时使有效带宽降低多达 50%。标准 FSDP 的所有集体通信都跨越全部进程，通信拓扑扁平，因此在扩展到数百个 GPU 时容易受到慢速节点间链路的严重拖累。HSDP 通过两层通信拓扑缓解了部分问题，但论文在实验中发现其学习性能可能出现明显退化。

联邦学习（FL）提供了一种互补思路：训练参与者可以以较低频率进行聚合，而不需要每个 batch 都做全量同步。论文没有把 FL 当作 DP 的替代品，而是把它作为一种更高层、更粗粒度的并行维度，叠加在 FSDP 和 HSDP 之上。

核心方法和技术细节

论文基于 NCCL 集体通信代价模型定义了几种方法的通信时间。设 $M$ 为层分片大小，$L$ 为层数，$S$ 为分片数，$R$ 为模型副本数，$F$ 为联邦组数，$A$ 为聚合频率。FSDP 每层需要两次 all-gather 和一次 reduce-scatter，其通信成本为：

$$C_{\text{FSDP}} = ML\left(3\frac{S-1}{S B_S}\right)$$

其中 $B_S$ 是 sharding group 内的最小带宽。HSDP 在 FSDP 基础上增加了跨 replication group 的梯度 all-reduce，成本为：

$$C_{\text{HSDP}} = ML\left(3\frac{S-1}{S B_S} + 2\frac{R-1}{R B_R}\right)$$

其中 $B_R$ 是 replication group 内的最小带宽。

联邦学习方式下，假设使用去中心化 FedAvg，每 $A$ 步对模型参数做一次 all-reduce，平均每 batch 通信成本为：

$$C_{\text{FL}} = ML\left[\frac{1}{A}\left(2\frac{F-1}{F B_F}\right)\right]$$

其中 $B_F$ 是联邦参与者之间的最小带宽。

论文提出的 FL+FSDP 把全局进程池分成 $F$ 个联邦组，每个联邦组内部运行一个完整的 FSDP，组内保持原有通信模式，每 $A$ 步在对应 shard 之间执行一次跨联邦组的 all-reduce。其通信成本为：

$$C_{\text{FL+FSDP}} = ML\left[3\frac{S-1}{S B_S} + \frac{1}{A}\left(2\frac{F-1}{F B_F}\right)\right]$$

FL+HSDP 则进一步利用 HSDP 内部的 replication 维度：每个联邦组运行一个 $R \times S$ 的 HSDP，跨联邦组只对每个 replication group 的一个副本做 all-reduce，再在组内广播给其他副本。其通信成本为：

$$C_{\text{FL+HSDP}} = ML\left[3\frac{S-1}{S B_S} + 2\frac{R-1}{R B_R} + \frac{1}{A}\left(2\frac{F-1}{F B_F} + \frac{1}{B_R}\right)\right]$$

该方法假设所有联邦组采用相同的分片策略，以保证不同联邦组之间 shard 一一对应。论文第 4.2 节指出，如果违反这一对齐约束，轻则引发运行时错误，重则导致语义错误的 shard 被平均，造成模型静默损坏。

创新点和贡献

论文有三个主要贡献。第一是混合通信方案：将 FSDP 或 HSDP 与 FedAvg 风格聚合交错，形成 FL+FSDP 和 FL+HSDP，把高频、重量级集体通信尽量限制在局部链路，跨联邦组的慢速通信仅以低频发生。第二是形式化通信代价分析：给出 FSDP、HSDP、纯 FL 以及两种新方法的平均每 batch 通信时间模型，并比较了在 $B_S \gg B_R \gg B_F$ 的层级带宽假设下各自的适用条件。第三是实验验证：在 128 节点、512 个 A100 GPU 的 Leonardo 超算上，用 Llama3.1 8B 预训练任务评估了吞吐和学习性能。

与 DiLoCo 和 Photon 等已有联邦大模型训练工作相比，本文的不同点在于：聚焦单集群内部的多级互连，而不是跨站点广域环境；FL 被用作 DP 之上的更高层并行原语，而不是替代 DP；通信拓扑显式建模 sharding、replication 和 federation 三个维度。

实验结果分析

实验在 CINECA 的 Leonardo 系统上进行，使用 128 个节点、共 512 个 A100 GPU，训练 Llama3.1 8B 约 150M 个 token，所有方法共享相同超参数，聚合频率 $A=8$。

表 3 显示，FSDP 训练时间最长，为 770.54 秒；HSDP 为 117.37 秒；FL+FSDP 为 95.87 秒；FL+HSDP 为 101.04 秒。FL+FSDP 处理相同数据比 FSDP 快 8.04 倍，比 HSDP 快 1.22 倍，比 FL+HSDP 快 1.05 倍。评估阶段，FL+FSDP 比 FSDP 快 8.48 倍。

学习性能方面，训练困惑度上 FSDP 为 496.00，HSDP 为 1151.47，FL+FSDP 为 403.29，FL+HSDP 为 254.31。评估困惑度分别为 486.64、1145.78、436.65、255.81。FL+HSDP 的评估困惑度比 FSDP、HSDP 和 FL+FSDP 分别低 1.83 倍、4.48 倍和 1.71 倍（表 3）。需要注意，这些数字仅在该数据集和实验设置下成立，不能推广为所有规模或模型下的普遍规律。

表 4 将训练时间拆分为前向、反向、优化器和聚合阶段。FSDP 的前向耗时 4.53 秒，反向 7.50 秒，是四种方法中最慢的。HSDP 和 FL 增强方法的前向时间约为 0.37 至 0.38 秒，差异很小。FL+FSDP 的反向时间为 0.89 秒，比 FSDP 快 8.43 倍，比 HSDP 快 1.61 倍。FL 增强方法额外增加聚合阶段，FL+FSDP 平均 0.17 秒，但标准差高达 0.55 秒，FL+HSDP 平均 0.14 秒，标准差 0.54 秒，说明聚合步骤耗时波动很大。

论文观察到 HSDP 的训练困惑度比 FSDP 更差，作者推测原因是 bfloat16 精度下浮点运算顺序差异被放大。对于 FL 增强方法的学习提升，作者推测与有效全局 batch size 有关：联邦组划分使全局 batch 只随联邦组数增长，而不是随进程总数增长，周期性 FedAvg 在全局层面稳定了收敛。

实践建议

对于在分层互连 HPC 集群上部署大规模分片数据并行训练的团队，这篇论文提供了一些可操作的工程启示。

首先，如果集群规模超过单个节点的 GPU 数量，不建议继续使用纯 FSDP 的扁平通信拓扑。在论文的实验设置下，FSDP 的前向和反向时间被跨节点集体通信严重拖慢，而 HSDP、FL+FSDP 和 FL+HSDP 通过层级通信显著降低了对慢速链路的依赖。若目标是极致吞吐，FL+FSDP 在表 3 中展示了最佳的数据处理速度；若目标是模型质量，FL+HSDP 在相同设置下取得了最低的评估困惑度。实际选择需要在吞吐和困惑度之间权衡。

其次，聚合频率 $A$ 是一个关键调参对象。$A$ 越小，跨联邦组通信越频繁，吞吐损失越大；$A$ 越大，组间模型发散时间越长，可能影响收敛。论文的公式显示，联邦聚合项被 $1/A$ 缩放，因此网络带宽差异越大，越可以适当增大 $A$ 来利用局部带宽。

第三，部署 FL+FSDP 或 FL+HSDP 时，必须保证所有联邦组采用相同分片策略，即 shard 数量和语义严格对齐。论文第 4.2 节特别警告，错位可能不会立即报错，而会导致静默模型污染。实践中应通过配置校验或断言机制避免这类风险。

第四，节点的物理分配应当尽量匹配逻辑拓扑。论文第 6 节指出，SLURM 的非确定性节点分配可能把联邦组拆到不同互连组，破坏 $B_S > B_R > B_F$ 的带宽假设，削弱 FL+HSDP 的优势。因此，建议在作业调度层面增加拓扑感知约束，优先让 sharding group 落在同一节点内，replication group 落在相邻节点或同一机架内，federation group 则在更远距离上形成。

最后，当网络层级更多或带宽差异更明显时，FL+HSDP 的三层拓扑更有价值；而在论文所测的集群中，FL+FSDP 已经达到了与 HSDP 相当的吞吐。因此，建议先测量目标系统的实际带宽层级，再决定是否需要引入 replication 维度的混合联邦方案，而不是盲目套用最复杂的拓扑。