借助联邦学习实现大规模分片数据并行加速

人工智能模型与高性能计算系统的协同扩展不断为二者的协同发展带来算法挑战。基础模型（FMs）就是一个典型例子：其训练通常需要数月时间，并使用数千块尖端 GPU。分片数据并行（DP）通过将数据和模型拆分到多块 GPU 上，是加速此类计算的主流策略。然而，在大规模部署时，尤其是在具有异构性能的多层互连网络上，该策略会产生难以承受的通信开销。受联邦学习（FL）高效通信原则的启发，本文提出了两种混合算法——FL+FSDP 和 FL+HSDP，将分片 DP 与 FedAvg 风格的聚合交错结合。这类方法将大规模 DP 部署解耦为更小、松耦合的联邦组，仅需极少的组间通信，同时使全局批量大小受联邦组规模限制。通信成本的形式化分析与实验验证证明了它们的可扩展性和灵活性。在 512 块 A100 GPU 上进行的 Llama3.1 8B 预训练表明，在超参数完全相同的情况下，FL+FSDP 和 FL+HSDP 的数据处理速度最高可提升 8.04 倍，评估困惑度最高可降低 4.48 倍，展现出更高的计算效率和更优的模型质量。这些特性源于通信开销的降低，以及全局批量大小相对于联邦组规模的受限增长。