Accelerating Sharded Data Parallelism at Scale with Federated Learning

📄 arXiv: 2609.20359v1 📥 PDF

作者: Gianluca Mittone, Marco Aldinucci

分类: cs.DC, cs.AI, cs.PF

发布日期: 2026-09-17

期刊: Euro-Par 2026: Parallel Processing - 32nd European Conference on Parallel and Distributed Processing, Pisa, Italy, August 24-28, 2026, Proceedings, Part II

DOI: 10.1007/978-3-032-35251-4_30


💡 一句话要点

提出FL+FSDP和FL+HSDP以解决大规模数据并行通信开销问题

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 分片数据并行 联邦学习 高性能计算 基础模型 通信效率 深度学习训练 模型优化

📋 核心要点

  1. 现有的分片数据并行方法在大规模部署时面临通信开销过大的挑战,影响训练效率。
  2. 本文提出FL+FSDP和FL+HSDP两种混合算法,通过将分片DP与联邦学习的聚合方式结合,降低通信成本。
  3. 实验结果表明,所提方法在数据处理速度和模型评估质量上均显著优于传统方法,提升幅度可达8.04倍和4.48倍。

📝 摘要(中文)

人工智能模型与高性能计算系统的协同扩展不断带来收敛的算法挑战。基础模型(FMs)需要在数千个先进GPU上进行数月的训练。分片数据并行(DP)是加速此类计算的主要策略,但在大规模部署时会产生巨大的通信开销。受联邦学习(FL)高效通信原则的启发,本文提出了两种混合算法——FL+FSDP和FL+HSDP,将分片DP与FedAvg风格的聚合交替进行。这些方法将大型DP部署解耦为较小的松耦合联邦组,减少了组间流量,同时保持全局批量大小受限于组的大小。通过正式的通信成本分析和实验验证,证明了其可扩展性和灵活性。实验结果显示,FL+FSDP和FL+HSDP在512个A100 GPU上进行Llama3.1 8B预训练时,数据处理速度提高了8.04倍,评估困惑度降低了4.48,展现出卓越的计算效率和模型质量。

🔬 方法详解

问题定义:本文旨在解决在大规模GPU集群中应用分片数据并行时,因通信开销过大导致的训练效率低下问题。现有方法在多层互连和异构性能环境下,难以有效协调数据和模型的分布,造成了显著的性能瓶颈。

核心思路:论文提出的FL+FSDP和FL+HSDP算法通过将分片DP与联邦学习的聚合机制相结合,形成松耦合的联邦组,从而减少组间的通信需求,同时保持全局批量大小的可控性。这样的设计旨在优化通信效率,提升整体训练速度。

技术框架:整体架构包括数据分片、模型分片和联邦学习聚合三个主要模块。首先将数据和模型分片到多个GPU上,然后在每个联邦组内进行局部训练,最后通过聚合机制整合各组的结果,形成全局模型更新。

关键创新:最重要的技术创新在于将联邦学习的高效通信原则引入到分片数据并行中,形成FL+FSDP和FL+HSDP两种新算法。这种方法有效地降低了通信开销,并通过松耦合的方式提高了系统的可扩展性。

关键设计:在算法实现中,关键参数包括联邦组的大小、局部批量大小和聚合频率等。损失函数设计上,保持了与传统DP一致性,同时在聚合过程中引入了FedAvg的思想,以确保模型更新的有效性。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

实验结果显示,FL+FSDP和FL+HSDP在512个A100 GPU上进行Llama3.1 8B预训练时,数据处理速度提高了8.04倍,评估困惑度降低了4.48,展现出显著的计算效率和模型质量提升,优于传统分片数据并行方法。

🎯 应用场景

该研究的潜在应用领域包括大规模深度学习训练、分布式计算系统和高性能计算环境。通过优化数据并行训练的通信效率,能够显著提升AI模型的训练速度和质量,推动更复杂模型的开发与应用,具有重要的实际价值和未来影响。

📄 摘要(原文)

The symbiotic scaling of artificial intelligence models and high-performance computing systems continually creates algorithmic challenges in their convergence. Foundation models (FMs) are a crucial example, requiring months-long training on thousands of cutting-edge GPUs. Sharded data parallelism (DP) is the dominant strategy to accelerate such computations by splitting data and models across multiple GPUs. However, it incurs prohibitive communication overhead when deployed at scale, particularly on multi-tier interconnects with heterogeneous performance. Inspired by the efficient communication principles of federated learning (FL), this work introduces two hybrid algorithms - FL+FSDP and FL+HSDP - interleaving sharded DP with FedAvg-style aggregations. Such approaches decouple large DP deployments into smaller, loosely-coupled federation groups, requiring minimal inter-group traffic while keeping the global batch size bounded by the groups' size. Formal analysis of communication costs and experimental validation prove their scalability and flexibility. A Llama3.1 8B pre-training on 512 A100 GPUs shows that, under identical hyperparameters, FL+FSDP and FL+HSDP achieve up to 8.04 faster data processing and 4.48 lower evaluation perplexity than their counterparts, demonstrating superior computational efficiency and improved model quality. These properties stem from reduced communication overhead and the bounded growth of the global batch size relative to the federation group size.