[Paper Review] The Future of Large Language Model Pre-training is Federated

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

Topic Modeling4 citations

TL;DR

This paper proposes a federated learning (FL) framework for pre-training billion-parameter large language models (LLMs) in a decentralized, collaborative manner, enabling data-rich institutions to train state-of-the-art models without centralized data centers. It demonstrates the first successful training of a billion-scale LLM in a heterogeneous FL setting, achieving performance competitive with centralized training while democratizing access to computational and data resources.

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.

Motivation & Objective

• To address the growing data and compute bottleneck in centralized LLM pre-training by leveraging underutilized global data and computational resources.
• To enable data-rich institutions—regardless of their compute or connectivity— to contribute to LLM pre-training without sharing raw data.
• To demonstrate that federated learning can scale to billion-parameter models with competitive performance to centralized training.
• To develop a reproducible, open-source system for federated LLM pre-training that supports heterogeneous clients and private data.
• To investigate whether larger federated models converge more easily than smaller ones, challenging conventional assumptions.

Proposed method

• Uses the Flower open-source FL framework to orchestrate decentralized training across multiple clients with varying hardware and data.
• Employs FedAvg with momentum-based aggregation to stabilize optimization and align client gradients across heterogeneous nodes.
• Applies local SGD with weight decay, clipping, and learning rate scheduling to reduce communication overhead and improve convergence.
• Implements a training recipe with fully transparent hyperparameters and configuration to ensure reproducibility.
• Uses server-side model validation and perplexity tracking to monitor convergence and detect overfitting.
• Enables model aggregation without direct data sharing, preserving privacy by design.

Experimental results

Research questions

• RQ1 Can federated learning achieve competitive LLM pre-training performance compared to centralized training, even with limited hardware and heterogeneous clients?
• RQ2 Does the scalability of federated training improve with increasing model size, contrary to expectations?
• RQ3 How do client and server model norms evolve during federated optimization, and what does this reveal about convergence dynamics?
• RQ4 To what extent does model averaging in FL regularize training and reduce noise in client gradients?
• RQ5 Can private, non-public data sources be effectively leveraged in LLM pre-training without centralizing data or violating privacy?

Key findings

• The first successful pre-training of a billion-scale LLM in a heterogeneous federated setting was achieved using only 1–8 GPU nodes from standard cloud providers.
• Federated training matched or exceeded centralized performance in terms of validation perplexity, with the server model serving as a soft upper bound on client spikes.
• Larger federated models (e.g., 350M parameters) found consensus across clients more easily than smaller models, defying initial expectations.
• The FedAvg pseudo-gradient norm decayed faster than local client gradients, indicating that averaging cancels conflicting or noisy optimization directions.
• The interplay between client and server models showed that the server model initially suppressed client norms but later stabilized and converged to the same norm, reflecting successful consensus formation.
• The system demonstrated that data-rich actors can become central players in LLM pre-training, even without massive compute infrastructure, by participating in a federated network.