Photon: Federated LLM Pre-Training
作者: Lorenzo Sani, Alex Iacob, Zeyu Cao, Royson Lee, Bill Marino, Yan Gao, Dongqi Cai, Zexi Li, Wanru Zhao, Xinchi Qiu, Nicholas D. Lane
分类: cs.LG, cs.DC
发布日期: 2024-11-05
备注: 13 pages, 9 appendix pages, 10 figures, 3 algorithms, 8 tables
💡 一句话要点
Photon：首个端到端联邦LLM预训练系统，实现低带宽下的全局规模模型训练。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 联邦学习 大型语言模型 预训练 分布式训练 数据隐私 低带宽 模型聚合
📋 核心要点
- 现有LLM训练受限于数据中心，高带宽需求阻碍了在弱连接GPU上进行更大规模的协同训练。
- Photon提出了一种基于联邦学习的端到端LLM训练系统，通过跨孤岛FL实现全局规模训练，并显著降低通信开销。
- 实验表明，Photon在训练7B模型时，困惑度优于集中式预训练，且训练时间随计算资源增加而减少，优于传统分布式训练。
📝 摘要（中文）
扩展大型语言模型（LLM）需要大量数据和计算资源，传统上，由于分布式训练的高带宽需求，这些资源被限制在数据中心内。像联邦学习（FL）这样的低带宽方法，如果能有效地用于预训练，就可以实现跨弱连接GPU的更大模型的协同训练。为此，我们推出了Photon，这是第一个完整的联邦端到端LLM训练系统，利用跨孤岛FL以最小的通信开销进行全球规模的训练。使用Photon，我们从头开始训练了第一个联邦解码器LLM系列。我们表明：（1）Photon可以以联邦方式训练高达7B的模型，同时达到比集中式预训练更好的困惑度；（2）Photon模型训练时间随着可用计算量的增加而减少，实现了与集中式相似的计算-时间权衡；（3）通过减少64倍-512倍的通信量，Photon的实际运行时间优于基线分布式训练方法35%。我们的方案对数据异构性具有鲁棒性，并且比以前的方法（如DiLoCo）收敛速度快两倍。这种惊人的数据效率源于一种独特的策略，即将小客户端批量大小与极高的学习率相结合，这得益于联邦平均对超参数的鲁棒性。因此，Photon代表了第一个经济的全球互联网范围LLM预训练系统。
🔬 方法详解
问题定义：论文旨在解决大规模语言模型（LLM）预训练过程中对高带宽和集中式数据中心的依赖问题。传统分布式训练方法需要大量数据传输，限制了在资源受限或数据分散的环境中训练大型模型。现有联邦学习方法在LLM预训练上的应用效果不佳，存在收敛速度慢、对数据异构性敏感等问题。
核心思路：Photon的核心思路是利用联邦学习的低带宽特性，实现跨多个数据孤岛（silos）的LLM预训练。通过在本地客户端进行模型训练，仅传输模型参数或梯度，从而显著降低通信开销。同时，论文提出了一种独特的训练策略，结合小批量大小和高学习率，以提高数据效率和收敛速度。
技术框架：Photon系统采用典型的联邦学习架构，包括以下主要模块：1）服务器：负责模型聚合和参数更新；2）客户端：在本地数据上进行模型训练；3）通信层：负责服务器和客户端之间的参数传输。训练流程包括：1）服务器初始化模型；2）服务器选择部分客户端参与训练；3）客户端下载模型参数，在本地数据上训练；4）客户端上传模型更新（例如梯度）；5）服务器聚合客户端的更新，更新全局模型；6）重复步骤2-5，直到模型收敛。
关键创新：Photon的关键创新在于：1）它是第一个完整的联邦端到端LLM预训练系统；2）它提出了一种结合小批量大小和高学习率的训练策略，显著提高了数据效率和收敛速度；3）它在实际运行时间上优于基线分布式训练方法，同时显著降低了通信量。
关键设计：Photon的关键设计包括：1）使用联邦平均算法进行模型聚合；2）采用小批量大小（具体数值未知）和高学习率（具体数值未知）进行本地训练；3）针对数据异构性，可能采用了某种正则化或梯度裁剪技术（具体细节未知）；4）模型架构采用decoder-only LLM，具体结构细节未知。
🖼️ 关键图片
📊 实验亮点
Photon在联邦学习环境下成功训练了高达7B参数的LLM，并且困惑度优于集中式预训练。通过减少64倍-512倍的通信量，Photon的实际运行时间比基线分布式训练方法提高了35%。此外，Photon的收敛速度是先前方法DiLoCo的两倍，显示出对数据异构性的更强鲁棒性。
🎯 应用场景
Photon为在数据隐私敏感或数据分散的环境中训练大型语言模型提供了新的可能性。它可应用于医疗、金融等领域，在保护用户隐私的前提下，利用分布在不同机构的数据进行模型训练，从而提升模型性能和服务质量。此外，Photon还可用于在资源受限的环境中训练LLM，例如在边缘设备或低带宽网络环境下。
📄 摘要（原文）
Scaling large language models (LLMs) demands extensive data and computing resources, which are traditionally constrained to data centers by the high-bandwidth requirements of distributed training. Low-bandwidth methods like federated learning (FL) could enable collaborative training of larger models across weakly-connected GPUs if they can effectively be used for pre-training. To achieve this, we introduce Photon, the first complete system for federated end-to-end LLM training, leveraging cross-silo FL for global-scale training with minimal communication overheads. Using Photon, we train the first federated family of decoder-only LLMs from scratch. We show that: (1) Photon can train model sizes up to 7B in a federated fashion while reaching an even better perplexity than centralized pre-training; (2) Photon model training time decreases with available compute, achieving a similar compute-time trade-off to centralized; and (3) Photon outperforms the wall-time of baseline distributed training methods by 35% via communicating 64x-512xless. Our proposal is robust to data heterogeneity and converges twice as fast as previous methods like DiLoCo. This surprising data efficiency stems from a unique approach combining small client batch sizes with extremely high learning rates, enabled by federated averaging's robustness to hyperparameters. Photon thus represents the first economical system for global internet-wide LLM pre-training.