PaperlayerView exact revision on arXiv

Source-linked AI summary

The Future of Large Language Model Pre-training is Federated

Lorenzo Sani, Alex Iacob, Zeyu Cao, Bill Marino, Yan Gao, Tomas Paulik, Wanru Zhao, William F. Shen, Preslav Aleksandrov, Xinchi Qiu, Nicholas D. Lane

arXiv:2405.10853v3cs.LGcs.AIcs.DC

TL;DR

LLM progress is limited by access to sufficient data and compute, much of which remains private or underused by centralized training. The paper presents Photon, a scalable federated pre-training system that enables institutions to collaborate without directly sharing private data. Photon supports billion-scale models up to 7B parameters, matches centralized performance with lower communication overhead, and remains robust to heterogeneity and partial participation.

  • Problem

    LLM improvement requires increasingly large datasets and compute, but centralized pre-training has limited access to private data and imposes high resource costs.

    Source: §1 · p. 2

  • Method

    Photon is a scalable federated pre-training system that coordinates heterogeneous data sources and compute providers without directly sharing private datasets.

    Source: §1 · p. 3, §4.1 · pp. 8–9

  • Results

    Federated training offers competitive performance with centralized training and far less communication overhead across model sizes up to 7B parameters.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Federated pre-training can mobilize distributed data and computational resources for collaborative billion-scale LLM training.

    Source: §1 · p. 3, §9 · p. 24

    Photon assumes fast and stable aggregator-to-node connectivity, with external communication negligible relative to computation for compute-intensive LLM training.

    Source: §4.3 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

Centralized LLM pre-training is constrained by scarce private data, expensive compute, and communication costs, motivating federated collaboration. The paper introduces Photon and reports competitive, scalable federated training through 7B parameters with resilience to heterogeneity and partial participation.

Motivation

  • Private data scarcity and rising compute costs threaten continued LLM scaling, while centralized training concentrates resources among well-resourced institutions.Public data may become insufficient, and centralized training requires costly, specialized hardware and long convergence runs.Source: §1 · p. 2
  • Federated learning can access previously untapped data and computation while avoiding direct sharing of private datasets.The proposed approach also reduces synchronization demands and can use moderately powerful hardware.Source: §1 · p. 2, Figure 1 · p. 2

Contribution

  • Photon is a practical, reproducible system for geographically distributed federated pre-training across heterogeneous devices, data sources, and trusted compute partnerships.It is built on Flower and supports affordable configurations using nodes with 1–8 GPUs.Source: §1 · p. 3

Results

  • Federated LLM training offers competitive performance with centralized training and far less communication overhead across model sizes up to 7B parameters.This result is reported as a broad comparison rather than for a single benchmark or task.Source: §1 · p. 3
  • Larger federated models require less frequent communication and receive a larger generalization boost than smaller models under Photon.The introduction also reports that overparameterized LLMs are less affected by statistical heterogeneity and partial participation than small FL models.Source: §1 · p. 3
  • Photon’s local execution engine balances computational loads and is reported to produce a faster, more robust pipeline than its centralized counterpart.The claim concerns the training pipeline enabled by Pollen.Source: §1 · p. 3

2 The Landscape of LLM Training

LLM pre-training combines demanding centralized optimization with growing data and hardware bottlenecks, while federated learning offers privacy-preserving collaboration but introduces heterogeneity challenges. Prior federated work has mainly addressed fine-tuning and parameter-efficient adaptation rather than full billion-scale pre-training.

2 The Landscape of LLM Training

  • Scaling laws motivate increasing both model and dataset sizes, but LLM pre-training requires massive corpora, large batches, long contexts, and extensive distributed computation.These requirements drive the use of data and model parallelism across many accelerators.Source: §2 · p. 3, §2.1 · pp. 3–4

2.1 Centralized Distributed Optimization

  • Distributed data parallelism replicates models and averages accumulated gradients, whereas model parallelism splits tensors across GPUs but becomes difficult to scale across machines.Ring AllReduce is sensitive to transmitted data and the slowest connection in the worker topology.Source: §2.1.1 · p. 4

2.1.2 Fully Sharded Data Parallelism

  • Fully sharded data parallelism reduces memory linearly with the data-parallelism degree but increases communication by 1.5× compared with standard DDP.It can be combined with model parallelism, activation checkpointing, and CPU offloading.Source: §2.1.2 · p. 4
  • Photon relaxes the hardware lower bound for participation by adjusting local computation and micro-batch sizes across organizations with independent resources.This allows weaker hardware to participate within reasonable efficiency limits.Source: §2.1.2 · pp. 4–5

2.1.3 Bottlenecks for generative pre-training of LLMs

  • Public language data may be exhausted, while scarce and costly accelerators must also be highly uniform and well connected for centralized synchronization.These constraints worsen with model size because model partitioning increases communication demands and optimization gaps.Source: §2.1.3 · p. 5

2.1.4 Mitigation of LLMs demands

  • Methods using pre-trained weights improve downstream access but do not resolve the bottleneck of who controls LLM pre-training and its data.The paper therefore motivates distributed pre-training that can use varied hardware and data, including training from scratch where feasible.Source: §2.1.4 · p. 5

2.2 Federated Learning and Local SGD

  • Federated learning keeps data local while clients train a shared model through repeated selection, local updates, aggregation, and global optimization.It can incorporate privacy mechanisms and reduce memorization and leakage by accessing previously unseen data.Source: §2.2 · pp. 5–6

2.3 Federated Fine-tuning and Parameter Efficient Fine-tuning of LLMs

  • Federated learning’s central challenges are data and systems heterogeneity, while prior LLM efforts largely focused on fine-tuning, PEFT, and prompt tuning rather than full pre-training.Federated PEFT can substantially reduce communication costs; federated LoRA is reported at 0.058% of full fine-tuning communication cost.Source: §2.2 · p. 6, §2.3 · pp. 6–7

3 Design Principles for Federated Generative Pre-Training of LLMs

The proposed federated pre-training design prioritizes inclusive access to data and computation, low communication requirements, hardware flexibility, and scalable local training. Photon connects data and compute providers while adapting execution to heterogeneous resources and model sizes.

Design principles

  • The federated pre-training paradigm is intended to create new models that can support later fine-tuning while reducing dependence on organizations willing to release open-source weights.Its design principles focus on data and hardware inclusivity, robustness, and efficiency.Source: §3 · p. 7

Broad Access to Data and Compute

  • Photon lets data-rich organizations participate without owning substantial compute by partnering with compute-rich, data-poor organizations.A client may contribute computing resources, data sources, or both.Source: §3 · p. 7

Limited Communication Requirements

  • Federated pre-training reduces synchronization demands so geographically distributed and poorly connected participants can contribute.The paper notes potential benefits for private data sources from regions with lower internet presence and for low-resource languages.Source: §3 · p. 7

Broad Hardware Inclusivity

  • Photon accommodates limited or poorly connected hardware by training locally across a client’s nodes, aggregating updates transparently, and adjusting round duration or local work.The supported range extends from a few medium-powered GPUs to small data centers.Source: §3 · p. 7

Scalable Local Training Pipelines

  • The system supports local pipelines from billion-parameter models with fully sharded data parallelism to simpler data parallel and CPU-offloaded training.This flexibility addresses participants with different model-size requirements and resource profiles.Source: §3 · p. 7

4 Photon Design

Photon is a federated pre-training system designed for private, heterogeneous data and compute across collaborating institutions. Its architecture and workflow support scalable local training, flexible participation, and compute-efficient communication.

4 Photon Design

  • Photon combines an aggregator, distributed LLM nodes, and private data sources to coordinate federated LLM pre-training across heterogeneous organizations.The system supports private data, limited communication, heterogeneous hardware, and collaborations between data holders and compute providers.Source: §4 · p. 8, §4.1 · pp. 8–10

4.1 Photon Architecture

  • The Photon Aggregator samples clients, broadcasts models and instructions, aggregates updates, supports secure communication, and checkpoints global state for recovery.It can use FedAvg, FedProx, or FedOPT, with asynchronous partial aggregation when the method permits it.Source: §4.1 · p. 8
  • Photon LLM Nodes adapt local training to available hardware, using distributed optimization, compression, checkpointing, or nested federation when connectivity is poor.Nodes may use DDP or FSDP on well-connected hardware and preserve local optimizer, data-loading, and epoch state privately.Source: §4.1 · p. 9, §4.2 · p. 11

4.2 Workflow Overview

  • Photon’s workflow couples siloed data sources with compute nodes, samples participants, performs local training and post-processing, aggregates updates, and evaluates models on private or public data.Public data can initialize or regularize training, while held-out local data supports distributed evaluation and personalization.Source: §4.2 · pp. 10–11

4.3 Communication-computation Scalability

  • For Photon’s cross-silo setting, computation dominates communication when clients perform sufficiently many local steps, unlike constrained cross-device FL.The described comparison uses 500 local steps and batch size 512 for Photon versus 4 steps and batch size 16 in prior work [85].Source: §4.3 · pp. 11–12
  • Federated local updates provide a practically larger, more varied batch, while unbiased sampling can reduce participation without hurting quality and may improve statistical efficiency 5–10×.For populations of roughly 32–128 clients, sampling 10% per round is described as potentially improving efficiency over full participation.Source: §4.3 · p. 12

5 Photon Implementation

Photon’s implementation uses modular microservices and adaptive execution to support reproducible federated LLM training. It also separates data production from computation through flexible streaming data sources.

5 Photon Implementation

  • Photon is implemented as microservices using Flower, PyTorch, Pollen, Hydra, and MinIO, with configuration and checkpointing designed to support reproducibility.The system stores large model-state binaries persistently and customizes Flower for checkpointing and reproducible sampling.Source: §5 · pp. 12–13

5.1 Parallel Execution within Photon LLM Nodes

  • Pollen adapts GPU collaboration to model size and hardware, supporting single-GPU, multi-GPU, and multi-machine execution through PyTorch Distributed.Photon selects DDP or FSDP when replicas or sharded models fit the available connectivity and memory constraints.Source: §5.1 · p. 13
  • When inter-machine connectivity is poor, Photon partitions data into a sub-federation, partially aggregates node models, and sends one client update to the server.This hierarchical strategy can also improve training efficiency within geographically distributed or poorly connected organizational data centers.Source: §5.1 · pp. 13–14

5.2 Streaming Data Sources

  • Photon supports heterogeneous data clients and flexible data-compute partnerships by letting compute nodes draw from arbitrary streamed data sources.Data can be streamed, cached, pre-tokenized, compressed, swapped between rounds, or newly incorporated without transferring large stored datasets.Source: §5.2 · p. 14

6 Methods

Photon is designed as a reproducible, scalable federated pre-training system that supports heterogeneous institutions, hardware, and private data. The evaluation uses controlled partitions, monitoring tools, reproducible open-source components, and model scales up to 7B parameters.

6 Methods

  • Photon provides a complete, scalable pipeline for experimentally evaluating federated LLM pre-training across heterogeneous clients and configurations.The system supports production-style orchestration, experimental modeling, and evaluation of whether federated learning can match centralized training.Source: §6 · p. 14

6.1 Reproducibility by Design

  • Reproducibility is supported through open-source components, disclosed training details, deterministic seeds, and reproducible client selection.The local pipeline uses MosaicML’s open-source MPT infrastructure, with every local training run and client-selection mechanism seeded.Source: §6.1 · p. 14

6.2 Experimental Tools

  • Photon monitors optimizer states, gradients, activations, client distances, cosine similarities, pseudo-gradients, and server momentum to detect federated divergence.It also supports automatic resumption, typed federated hyperparameters, resource cleanup, and per-device micro-batch selection through VRAM-aware search.Source: §6.2 · p. 15

6.2.1 Data Partitioning

  • Heterogeneous datasets are partitioned by meaningful categories and disjoint buckets, enabling complex cross-silo topologies without additional runtime bookkeeping.The partitions model differences such as source or genre in The Pile and language in mC4, while ensuring clients draw from disjoint subsets.Source: §6.2.1 · p. 15

6.3 Experimental Design

  • The evaluation tests homogeneous IID data, naturally heterogeneous institutional data, partial participation, and model scaling to assess federated optimization’s main adoption challenges.Homogeneous experiments use randomly sharded C4, while heterogeneous experiments use semantically distinct sources from The Pile.Source: §6.3 · pp. 15–16

6.4 Model Scaling

  • The analysis uses Hoffmann et al. [6] scaling laws, adopting 20 tokens per parameter, adjusting for a 50K rather than about 32K vocabulary, and comparing recommendations with MosaicML.All experiments use 500 local steps per round, executed by every federation client.Source: §6.4 · p. 16, Table 1 · p. 16, §6.5 · p. 17

6.5 Training Setup

  • Training uses heterogeneous A40, A100, and H100 servers across countries, with local AdamW, server FedMom, a 50,368-token tokenizer, and reported federated hyperparameters.The cross-silo setting relies on federated learning’s relaxed synchronization requirements to coordinate geographically distributed hardware.Source: §6.5 · p. 17, Table 3 · p. 17

7 Evaluation

Photon’s evaluation finds federated LLM training viable across increasing model scales and naturally heterogeneous data. Federated models approach centralized performance, while aggregation stabilizes training and limits activation growth.

7.1 Federated Optimization is Effective and Stable across Scales:

  • For 75M, 125M, 350M, and 1.3B models, federated training becomes more stable with scale, with centralized and federated performance near-identical at 1.3B.Aggregation initially produces client-perplexity spikes, then reaches consensus more quickly for larger models and provides a regularizing effect.Source: Figure 3 · p. 18, §7.1 · p. 18
  • Federated optimization reaches comparable centralized performance and remains viable from small models through 7B parameters.The evaluation reports that performance gaps shrink with model size, continuing from 1.3B to 3B and 7B models.Source: §7 · p. 17, §7.1 · p. 18

7.2 Federated Optimization is Robust to Heterogeneity:

  • Naturally heterogeneous The Pile partitions converge robustly: despite greater client-loss variance, server validation loss follows a trend comparable to IID training.The reported absolute performance changes from 50 to 35 for 75M and from 50 to 30 for 125M because the datasets differ.Source: §7.2 · p. 19, Figure 4 · p. 19
  • Federated aggregation limits activation growth under natural heterogeneity, while centralized activations increase rapidly and may accompany divergence.The authors attribute periodic reductions in activation norms to aggregation-injected noise and report recovery from temporary validation-accuracy decreases.Source: §7.2 · p. 19, Figure 4 · p. 19

7.3 Larger Models Improve Consensus:

Larger federated models reach client consensus more quickly and train more stably. Aggregation shifts from early turbulence to a regularizing effect at convergence.

7.3 Larger Models Improve Consensus:

  • The 1.3B model reaches consensus after only 4 rounds, whereas the 75M model requires over 20 rounds.The 75M model continues to show occasional convergence turbulence, while the 350M model achieves full client-model alignment.Source: §7.3 · p. 20
  • As model size increases, federated stability improves and final performance approaches centralized training, becoming almost identical for the 1.3B model.Aggregation raises local perplexity during the transient phase but lowers it at convergence through a regularizing effect.Source: §7.3 · pp. 19–20, Figure 5 · p. 20

7.4 Federated Optimization is Robust to Partial Participation:

Federated training remains effective with sparse client participation and exhibits convergent client-server dynamics. Partial participation reaches full-participation performance while reducing parallel compute, and client updates align as training proceeds.

7.4 Federated Optimization is Robust to Partial Participation:

  • Sampling 4 of 64 clients per round, or 6.25% participation, reaches the same performance as full participation while using half the parallel compute.The resulting robustness can reduce hardware utilization and support simultaneous federated workloads.Source: §7.4 · p. 20, Figure 6 · p. 21

7.5 Interplay of Client and Server Models:

  • Aggregation initially pulls back client-model norms, then increases the averaged-client norm until server and client models converge.The same norm dynamics appear under heterogeneous data and partial participation.Source: §7.5 · p. 20, Figure 7 · p. 21

7.6 Federated Optimization Aligns Client Gradients:

  • The average client pseudo-gradient begins much larger than local gradients but becomes similar or smaller as client models converge.The same gradient-alignment trend holds for heterogeneous-data and partial-participation experiments.Source: §7.6 · pp. 20–21

7.7 Pushing the Boundary on Federated LLM Size

Federated pre-training remains competitive as model size reaches 7B, with larger models outperforming centralized counterparts on perplexity. Among outer-optimizer choices, standard FedAvg is most robust and achieves the lowest final perplexity.

7.7 Pushing the Boundary on Federated LLM Size

  • Larger 3B and 7B federated models consistently outperform centralized counterparts on client-training and server-validation perplexity.The validation-perplexity gap is larger than for the 1.3B model, suggesting continued improvement in federated optimization.Source: §7.7 · p. 21, Figure 9 · p. 22

7.8 On the choice of the outer optimizer

  • The experiments compare AdamW locally with FedAvg or server-side Nesterov momentum and examine batch sizes matched to centralized training.The batch-size analysis considers both the centralized local batch size and a reduced local size whose effective aggregate matches centralized pre-training.Source: §7.8 · pp. 21–22
  • Standard FedAvg achieves the lowest final perplexity and is most robust to divergence across large- and small-batch settings.Server-side momentum and preserved local optimizer states can provide initial speedups but fail to approach centralized-training perplexity.Source: §7.8 · p. 22, Figure 10 · p. 23
  • Using stateless clients supports partial participation, dropout tolerance, and new federation entrants without staleness.Synchronizing optimizer states can increase communication costs threefold, while unsynchronized states can increase client drift under data heterogeneity.Source: §7.8 · p. 22
  • Server-side momentum and retained local optimizer states increase model l2 norms, causing eventual divergence in cross-entropy.Figure 10 compares standard FedAvg with SGD+N and FedAvg-KeepOpt under both batch-size conditions.Source: Figure 10 · p. 23

7.9 Evaluation Harness

Photon models scale in downstream in-context learning performance across model sizes. The largest model wins most reported comparisons across the two evaluation groups.

7.9 Evaluation Harness

  • The largest Photon model wins 10 of 14 in-context learning comparisons across Tables 5 and 6.This demonstrates downstream utility despite using a pre-training dataset not optimized for downstream performance.Source: §7.9 · p. 23, Table 6 · p. 23
  • The largest model wins 5 of 6 comparisons in the second in-context learning evaluation group.Source: Table 6 · p. 23

8 Future Work

Future work will extend Photon’s scale, evaluate broader downstream capabilities, and investigate how federated data heterogeneity and privacy affect model development. The authors also plan to reduce remaining runtime overheads to broaden participation.

8 Future Work

  • The authors plan to fine-tune federated models on established benchmark tasks and scale both federation population and model size.Source: §8 · p. 24
  • Future studies will examine how languages, genres, larger private datasets, and federated data heterogeneity affect model capabilities.The work will compare the potential advantages of collaboration with the isolated centralized paradigm.Source: §8 · p. 24
  • The authors will analyze privacy-preservation strengths and weaknesses and develop optimizations for the remaining execution-runtime overheads.These efforts aim to support broader participation in federated learning.Source: §8 · p. 24

9 Conclusions

The paper demonstrates billion-scale federated generative pre-training in a heterogeneous setting, with robustness to data heterogeneity and partial participation. Photon provides a collaborative, reproducible, and scalable path for pooling distributed computational and data resources.

9 Conclusions

  • Photon produces the first federated billion-scale model fully trained in a heterogeneous federated-learning setting.Source: §9 · p. 24
  • Federated training is robust to both data heterogeneity and partial participation.Source: §9 · p. 24
  • The complete system enables collaborative, reproducible, and scalable LLM pre-training while pooling computational and data resources distributed across the planet.The system is built on an open-source framework and the training recipe is fully disclosed for future development and collaboration.Source: §9 · p. 24

A Additional Results

Additional experiments show that server momentum and decaying learning rates enforce consensus despite heterogeneous data, while convergence trends remain unchanged under 6% client participation. Gradient analyses further distinguish heterogeneity-driven pseudo-gradient decay from learning-rate effects across model sizes.

Optimization dynamics

  • Server-side Nesterov momentum with β = 0.7 tracks an exponential moving average of the global model across 75M to 1.3B experiments.The corresponding global-model l2 norms are shown for four model sizes.Source: Figure 11 · p. 37

Statistical heterogeneity

  • Server momentum and a decaying learning rate enforce consensus between global and client models on heterogeneous The Pile partitions.The figure reports this behavior for 75M and 125M experiments.Source: Figure 12 · p. 37

Partial participation

  • With 4 of 64 clients per round, the global model follows the same trend as under full participation.This partial-participation setting samples only 6% of the federated population.Source: Figure 13 · p. 38

Statistical heterogeneity

  • On heterogeneous data, faster FedAvg pseudo-gradient decay than local-gradient decay indicates adaptation to data heterogeneity rather than learning-rate decay.This comparison uses 75M and 125M The Pile experiments.Source: Figure 14 · p. 38

Partial participation

  • Under partial participation, the relationship between FedAvg pseudo-gradients and local steps matches the full-participation case.The comparison uses 75M and 125M C4 experiments with a small client sample each round.Source: Figure 15 · p. 38
Loading 2405.10853v3…