[论文解读] The Future of Large Language Model Pre-training is Federated

Lorenzo Sani, Alex Iacob+9 more|arXiv (Cornell University)|May 17, 2024

Topic Modeling被引用 4

一句话总结

本文提出了一种联邦学习（FL）框架，用于以去中心化、协作的方式预训练千亿参数规模的大语言模型（LLMs），使数据丰富的机构能够在不依赖集中式数据中心的情况下训练最先进水平的模型。该工作首次在异构联邦学习设置下成功训练了千亿规模的LLM，在性能上与集中式训练相当，同时实现了计算和数据资源的民主化。

ABSTRACT

Generative pre-trained large language models (LLMs) have demonstrated impressive performance over a wide range of tasks, thanks to the unprecedented amount of data they have been trained on. As established scaling laws indicate, LLMs' future performance improvement depends on the amount of computing and data sources they can leverage for pre-training. Federated learning (FL) has the potential to unleash the majority of the planet's data and computational resources, which are underutilized by the data-center-focused training methodology of current LLM practice. Our work presents a robust, flexible, reproducible FL approach that enables large-scale collaboration across institutions to train LLMs. We propose a scalable deployment system called Photon to enable the investigation and development of this new training paradigm for LLM pre-training. We show that Photon can be used by organizations interested in collaborating with their private data sources and computational resources for pre-training LLMs with billions of parameters. This paradigm would mobilize more computational and data resources while matching or potentially exceeding centralized performance. We further show the effectiveness of the federated training scales with model size and present our approach for training billion-scale federated LLMs using limited resources. Thus far, we have used Photon to train LLM models to the size of 7B parameters and anticipate larger models being completed in the near future. Finally, we show that LLM training is highly resilient to the classical challenges of federated statistical and hardware heterogeneity. Furthermore, we show that convergence is robust to partial participation, opening the avenue for compute-efficient collaborative training. Photon will help data-rich actors to become the protagonists of LLMs pre-training instead of leaving the stage to compute-rich actors alone.

研究动机与目标

• 为解决集中式LLM预训练中日益增长的数据与计算瓶颈，通过利用全球范围内未被充分利用的数据和计算资源。
• 使数据丰富的机构——无论其计算能力或网络连接如何——都能在不共享原始数据的前提下参与LLM预训练。
• 证明联邦学习可扩展至千亿参数规模的模型，并实现与集中式训练相当的性能。
• 开发一个可复现的、开源的联邦LLM预训练系统，支持异构客户端和私有数据。
• 探究更大的联邦模型是否比更小的模型更容易收敛，从而挑战传统假设。

提出的方法

• 使用Flower开源联邦学习框架，在具有不同硬件配置的多个客户端之间协调去中心化训练。
• 采用带有动量的FedAvg聚合方法，以稳定优化过程并对齐异构节点上的客户端梯度。
• 应用本地SGD配合权重衰减、梯度裁剪和学习率调度，以减少通信开销并提升收敛性。
• 实施一套完整的训练方案，包含完全透明的超参数与配置，以确保可复现性。
• 通过服务端模型验证和困惑度追踪来监控收敛过程并检测过拟合。
• 实现模型聚合而无需直接共享数据，从设计上保障隐私。

实验结果

研究问题

• RQ1 在硬件资源有限且客户端异构的情况下，联邦学习是否仍能实现与集中式训练相当的LLM预训练性能？
• RQ2 联邦训练的可扩展性是否随着模型规模增大而提升，与预期相反？
• RQ3 在联邦优化过程中，客户端与服务端模型范数如何演变？这揭示了怎样的收敛动力学特征？
• RQ4 联邦学习中的模型平均在多大程度上起到正则化作用，从而减少客户端梯度中的噪声？
• RQ5 能否在不集中化数据或违反隐私的前提下，有效利用私有的、非公开的数据源进行LLM预训练？

主要发现

• 仅使用标准云服务提供商提供的1至8个GPU节点，首次成功实现了在异构联邦设置下的千亿规模LLM预训练。
• 联邦训练在验证困惑度方面达到或超过集中式训练性能，服务端模型作为客户端峰值的软上界。
• 更大的联邦模型（如3.5亿参数）比更小的模型更容易在客户端之间达成共识，这与初始预期相反。
• FedAvg伪梯度范数的衰减速度超过本地客户端梯度，表明平均过程能有效抵消冲突或噪声的优化方向。
• 客户端与服务端模型之间的相互作用表明，服务端模型最初抑制了客户端范数，但随后稳定并收敛至相同的范数值，反映出成功的共识形成过程。
• 该系统表明，数据丰富的参与者即使缺乏大规模计算基础设施，也能通过参与联邦网络成为LLM预训练的核心力量。