Photon: Federated LLM Pre-Training

Photon: Federated LLM Pre-Training

Speakers: Lorenzo Sani, Alex Iacob, Zeyu Cao, Royson Lee, Nicholas D. Lane

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This article delves into Photon, a pioneering system designed for Federated LLM pre-training, as presented by Lorenzo Sani and his collaborators from Flower Labs, the Machine Learning Systems Group at the University of Cambridge, and visiting researchers from Bup and Zhejiang University. The talk outlines a radical departure from conventional large language model (LLM) training paradigms, proposing a distributed, internet-scale approach to address the escalating challenges of compute infrastructure, training robustness, and data accessibility. Photon aims to leverage globally distributed computing resources and decentralized data sources, moving away from monolithic, tightly coupled data centers.

The core motivation behind Photon is to tackle three critical bottlenecks hindering the continued scaling of LLM pre-training: the immense and costly compute requirements, the fragility of current distributed training setups, and the impending scarcity of publicly available training data. By adopting a federated learning framework, Photon seeks to unlock vast, untapped resources and capabilities, making large-scale LLM development more sustainable, resilient, and inclusive. The system’s design emphasizes flexibility, fault tolerance, and communication efficiency, demonstrating that high-quality LLM pre-training is not only feasible but also highly advantageous in a federated environment.

Ultimately, Photon represents a significant step towards democratizing access to large-scale AI model development. By enabling training across heterogeneous, geographically dispersed nodes, it opens up possibilities for organizations and researchers to contribute to and benefit from LLM advancements without needing to centralize massive datasets or invest in prohibitive supercomputer clusters. The presented results showcase Photon's ability to achieve comparable or superior model quality, substantial training speedups, and enhanced robustness, underscoring its potential to reshape the future of LLM pre-training.

Background

▶ Watch: Introduction to Photon: Federated LLM Pre-training (0:00)

The landscape of large language model (LLM) pre-training is currently dominated by an insatiable demand for computational resources, vast datasets, and highly specialized infrastructure. As models like GPT-3, PaLM, and Llama continue to push the boundaries of scale and capability, three major bottlenecks have emerged, threatening the sustainability and accessibility of this critical research and development area. Photon was conceived as a direct response to these systemic challenges, aiming to provide a fundamentally different paradigm for LLM pre-training.

The first and perhaps most visible bottleneck is compute. Training ever-larger LLMs demands immense computing infrastructure, typically relying on massive GPU clusters interconnected with high-speed, low-latency networks. Building and maintaining these supercomputing data centers is incredibly expensive, environmentally impactful, and increasingly impractical to scale further. The talk highlights this by noting that the standard approach involves bigger and bigger GPU clusters, but the cost and physical limitations are becoming prohibitive. This concentration of resources also creates a single point of failure and limits participation to only the wealthiest or most resource-rich organizations.

Secondly, robustness presents a significant headache in large-scale training. Current LLM training jobs are often tightly coupled, meaning the failure of even a single node can crash the entire training process, leading to costly restarts and substantial delays. A compelling example cited is Meta's Llama 3 training, which experienced "dozens of interruptions" over a 54-day period, with a staggering 78% of these failures attributed to unpredictable hardware issues. Such fragility underscores the need for more resilient training methodologies that can gracefully handle node failures without derailing months of computational effort.

The third critical challenge revolves around data. Large models require enormous text corpora, and the rate at which these models are growing is quickly outstripping the supply of publicly available text data. An analysis from Villa predicts that the world could "run out of publicly available text data on the internet by next year" if current scaling trends persist. This impending data scarcity necessitates new strategies for accessing and utilizing data beyond mere web scraping. Federated learning emerges as a promising solution here, offering a mechanism to tap into decentralized data sources, such as private organizational datasets, that cannot be readily centralized due to privacy, proprietary, or logistical constraints. By enabling training on distributed data, federated learning could unlock a vast, currently inaccessible pool of information, ensuring the continued growth and improvement of LLMs.

These profound challenges – unsustainable compute scaling, brittle training processes, and data scarcity – collectively motivated the development of Photon. The project's core idea is to leverage the distributed, internet-wide nature of federated learning to transform these bottlenecks into opportunities, fostering a more robust, efficient, and data-rich environment for LLM pre-training.

Key Findings

▶ Watch: Photon Design Principles for Federated LLM (4:25)

Photon’s development and evaluation yielded several critical findings that underscore the viability and advantages of federated LLM pre-training:

• First System for Federated LLM Pre-Training: Photon is presented as the first system capable of pre-training large language models in a federated manner. This demonstrates that such a paradigm shift is not only theoretically possible but also practical and effective across various federated settings, including those with widely distributed and heterogeneous resources.
• Comparable or Superior Model Quality: A crucial finding is that federated training with Photon does not compromise model performance or quality. In fact, for a fixed amount of training time, Photon-trained models were found to be "as good or better" than those trained using centralized approaches. This is a significant validation, indicating that the benefits of federated learning do not come at the cost of model efficacy.
• Significant Training Time Reduction: Photon can substantially reduce the overall training time, particularly in environments with low network bandwidth. This efficiency gain stems from its ability to minimize communication overheads compared to traditional centralized data parallelism.
• Enhanced Resource Utilization: The system achieves higher GPU utilization rates, observed to be "slightly higher" and reaching up to 90% compared to centralized counterparts. This indicates that Photon clients are performing more total compute or work on data, making better use of available hardware.
• 2X Throughput Improvement for Large Models: For a 13 billion parameter model trained on 32 H100 GPUs across four nodes, Photon demonstrated a striking 2X higher throughput compared to a centralized approach. This impressive gain is directly attributed to its communication efficiency.
• Unimpacted Convergence with Reduced Communication: Despite significantly reducing the frequency of communication between nodes, Photon ensures that model convergence is not adversely affected. The aggregation steps effectively improve model performance without introducing the typical perplexity spikes often seen at round boundaries in other federated literature.
• Robustness to Client Failures: Photon significantly enhances training robustness. Experiments show that training can continue "seamlessly" and converge without particular intervention even when clients systematically fail. This contrasts sharply with the fragility of traditional tightly-coupled distributed training.
• Enables Internet-Wide Scale-Out: The system’s design principles, particularly its handling of hardware heterogeneity and stateless client design, make it uniquely suited for scaling out LLM training across diverse, internet-connected nodes, utilizing existing computing systems spread globally.
• High Communication Efficiency: A core contribution is Photon’s inherent communication efficiency. By performing extensive local work before aggregating updates, it drastically cuts down on network traffic, which is especially beneficial in bandwidth-constrained cross-silo environments.

These findings collectively highlight Photon's potential to address the core challenges of LLM pre-training, offering a path towards more efficient, resilient, and broadly accessible large-scale AI development.

Technical Deep Dive

▶ Watch: Achieving Federated Robustness with Stateless Client Design (6:40)

The technical prowess of Photon lies in its innovative design principles and sophisticated mechanisms for handling the inherent complexities of a federated cross-silo environment. The system was engineered from the ground up to overcome the limitations of traditional LLM training, focusing on data decoupling, flexible compute scaling, and robust fault tolerance.

Photon Design Principles

1. Access to Data: A cornerstone of Photon's design is the decoupling of data sources from the computing workers. Unlike conventional setups where data must reside locally with the compute, Photon allows clients to stream data from various sources—be it private collaborators, external providers, or public datasets. This flexibility is crucial for unlocking training on datasets that are distributed across different organizations or geographical locations, circumventing the need for data centralization which is often impractical or prohibited by privacy regulations. This principle directly addresses the impending data scarcity by making vast, decentralized data pools accessible for LLM pre-training.

1. Flexible Compute Scaling (Scale-Out Computing Strategy): Photon embraces a scale-out computing strategy, moving away from reliance on single, monolithic superclusters. Instead, it can harness a multitude of distributed nodes over the internet, effectively utilizing whatever computing resources are available, wherever they may be. This approach facilitates easy expansion to new participants, even if they introduce heterogeneity in computing resources. The ability to leverage "GPUs in the wild" transforms a potential limitation into a strength, offering a cost-effective and globally scalable compute infrastructure.

1. Federated Robustness (Fault Tolerance via Stateless Local Design): Critical for real-world deployment, Photon's training pipeline is designed to be stateless locally on the clients. This means that clients do not retain any essential training state between rounds. All crucial and relevant information, such as the latest model parameters and optimizer states, is managed and stored on the server through checkpointing and the adopted optimization procedure. Consequently, if a client fails or drops out, the entire workload does not stop, as the server retains all necessary information. This design significantly enhances fault tolerance, making the training process remarkably resilient against unpredictable hardware failures, a common issue in large-scale distributed systems.

Handling Heterogeneity in Dynamic Environments

A significant practical challenge in federated environments is the heterogeneity of client computing resources. Photon addresses this dynamically and intelligently:

• Strongly Connected GPUs: If a client possesses multiple GPUs that are highly connected (e.g., via NVLink, Infiniband, or DMA), Photon treats them as a single logical client. Within this logical client, these GPUs can effectively utilize advanced parallelism techniques such as tensor parallelism or pipeline parallelism to accelerate local computation.
• Weakly Connected GPUs: Conversely, if a client has GPUs that are weakly connected (e.g., through the internet or located in different distant nodes), employing high-bandwidth parallelism would introduce communication bottlenecks and slow down the process. In such cases, Photon treats these resources as a "mini federation." Each individual computing resource runs separate workloads, and their results are aggregated locally before being sent back to the central server.
• Dynamic Adaptation: Crucially, Photon makes these decisions "on the fly." At every training round, for every client, the system can adapt its local configuration based on changes in hardware setup or the machine learning pipeline requirements.
• Automatic Batch Size Tuning: A key enabler for this dynamic heterogeneity support is the automatic tuning of the batch size per device. This, combined with specific machine learning insights that allow for smaller batch sizes and higher learning rates compared to centralized training, enables a much more efficient distribution of workloads, especially compared to standard data parallelism approaches. This intelligent resource allocation is fundamental to enabling internet-wide training utilizing diverse existing systems.

Communication Efficiency and Optimization

Photon’s remarkable communication efficiency is a core technical advantage. In a federated learning setup, clients perform a significant amount of "local work" before synchronizing with the server. They compute partial gradients and apply them over many local steps (e.g., 1,024 steps for the 13B model). Only after these numerous local updates do they share the updated model parameters with the server for aggregation. This drastically reduces the frequency of communication compared to traditional data parallelism, where gradients are typically averaged across the entire cluster for each mini-batch. This reduction in communication overhead is particularly impactful in low-bandwidth environments, contributing to the observed 2X throughput gains.

For the 13B parameter model, the research team adopted a "new optimizer" compared to the one used in the main paper, indicating ongoing research into further enhancing convergence and performance within the federated paradigm. The convergence curves clearly illustrate how model performance improves at the round boundaries due to aggregation, notably without the typical perplexity spikes that can occur when clients readapt to local datasets. This suggests a well-tuned balance between local computation and global synchronization.

Model Architectures and Future Directions

The experiments primarily utilized dense decoder-only models. Acknowledging a current limitation, the speakers noted that clients need to be able to hold the full model in memory to participate. Future work aims to broaden participation by enabling even smaller clients to contribute, potentially through exploring different model architectures that are less dense and more sparse, or by implementing techniques like model parallelism across client nodes. The team is also actively developing new optimization procedures to scale to even larger models (e.g., 100 billion parameters) by mitigating the communication overhead associated with transferring massive parameter chunks.

Experimental Setup & Results

▶ Watch: Adaptive Resource Management via Mini-Federations (8:00)

To validate the capabilities of Photon, a series of rigorous experiments were conducted, comparing its performance against standard centralized data parallelism techniques across various scales and network conditions.

Datasets and Model Sizes

Initial evaluations involved training language models up to billion-size parameters from scratch. Standard benchmark datasets were employed, including the widely used C4 data set, The Pile, and a new partitioning of the FineWeb dataset. These diverse corpora ensured a comprehensive assessment of Photon's ability to handle different data distributions and scales.

For more advanced testing and pushing the boundaries for MLSys, the team recently scaled up to a 13 billion parameter model. This represented a significant increase in complexity and resource demand, allowing for a more direct comparison with state-of-the-art LLM training efforts. The models used were consistently dense decoder-only models.

Hardware and Network Environment

The initial experiments were designed to showcase Photon's ability to operate across "many modest nodes" distributed over the internet, emphasizing a "scale-out computing strategy." This implied leveraging heterogeneous hardware and potentially lower-bandwidth connections typical of real-world distributed environments.

For the 13 billion parameter model experiments, a more specific setup was used:

• Hardware: 32 H100 GPUs distributed across four nodes.
• Network: These four nodes were initially located in the same physical data center but across different sub-clusters. Crucially, the peer-to-peer network speed between these nodes was approximately 3 gigabit per second (Gbps). This bandwidth is considerably lower than what would be expected in a tightly coupled, high-speed cluster (e.g., using Infiniband), simulating a more realistic cross-silo or wide-area network environment.
• Baseline: Each Photon client in this setup corresponded to a data parallel worker in the centralized standard approach, ensuring a fair comparison.

Key Metrics and Headline Numbers

The evaluation focused on several key performance indicators: model quality, training time, throughput, and resource utilization.

1. Model Quality and Performance:

• Result: Photon-trained models achieved quality "as good or better" than centralized models for a fixed training duration. This indicates that the federated approach does not degrade the ultimate performance of the LLM.
• Convergence: The convergence curves for the 13B model showed that performance consistently improved at round boundaries due to aggregation, without the typical "spike in perplexity" or readaptation issues often seen in other federated learning literature. This suggests the new optimizer and communication strategy are highly effective.

1. Training Speed and Throughput:

• Result: Photon significantly reduced overall training time, particularly in low-bandwidth environments, by cutting down communication overheads.
• Headline Number: For the 13B parameter model, Photon achieved an impressive 2X higher throughput compared to the centralized data parallelism baseline. This substantial speedup is attributed to its highly communication-efficient design.

1. Resource Utilization:

• Result: Photon clients demonstrated higher GPU utilization.
• Headline Number: GPU utilization was observed to be "slightly higher" than centralized counterparts, reaching up to 90%. This indicates that the distributed clients are performing more useful work, making better use of the available computational resources.

1. Robustness:

• Result: The system exhibited remarkable resilience to client failures. New experiments detailed in the paper (though not explicitly shown in the talk for time) demonstrated that training "doesn't need to stop at any time" and "can converge seamlessly" even with systematic client failures. This is a crucial advantage over fragile centralized systems.

The experimental results definitively confirm that Photon can efficiently pre-train large language models with high quality in a federated manner, demonstrating significant advantages in throughput, resource utilization, and robustness, especially when operating over typical internet-scale network conditions rather than ultra-high-speed data center interconnects.

Practical Implications

▶ Watch: Photon Evaluation and Key Results Overview (9:30)

The development of Photon carries profound practical implications for various stakeholders involved in the lifecycle of large language models, from model builders to infrastructure teams and deployers. Its federated approach addresses critical pain points in traditional LLM training, offering a more sustainable and accessible path forward.

For Practitioners and Model Builders

• Access to Decentralized Data: Photon directly tackles the looming data scarcity issue by enabling training on decentralized data sources. This means model builders can now leverage private, proprietary, or geographically distributed datasets that cannot be centralized due to privacy concerns, regulatory compliance (e.g., GDPR, HIPAA), or logistical challenges. This unlocks a vast, untapped reservoir of information, allowing for the creation of more specialized, diverse, and robust LLMs.
• Broader Collaboration: The ability to train models across organizational boundaries without sharing raw data fosters new avenues for collaboration. Multiple institutions or companies can collectively contribute to a shared model while keeping their sensitive data localized, accelerating research and development in fields requiring data silos.
• Reduced Barrier to Entry for LLM Pre-training: By making efficient use of widely distributed and potentially heterogeneous compute resources, Photon lowers the barrier to entry for pre-training large models. Smaller organizations or academic institutions without access to multi-million dollar superclusters can now participate in LLM development by pooling more modest resources.

For Infrastructure Teams and Deployers

• Cost-Effective Compute Scaling: Instead of investing in incredibly expensive, ever-larger centralized GPU clusters, infrastructure teams can now leverage existing, distributed computing systems, including "GPUs in the wild." This scale-out computing strategy offers a more economical and flexible way to expand compute capacity.
• Enhanced System Robustness: The stateless local design of Photon clients and the server-side checkpointing mechanism provide unparalleled fault tolerance. Infrastructure teams will experience significantly fewer costly training interruptions due to single-node failures. This translates to higher uptime, reduced operational overhead, and faster development cycles, addressing the prevalent issue of fragility in current large-scale distributed training.
• Efficient Wide-Area Network Utilization: Photon’s high communication efficiency is a game-changer for deploying training across geographically dispersed nodes connected by typical internet bandwidth. The 2X throughput gains observed even with 3 gigabit per second peer-to-peer speeds demonstrate that large models can be trained effectively over wide-area networks, reducing the need for specialized, expensive, low-latency interconnects.
• Heterogeneity Management: The system’s dynamic adaptation to heterogeneous hardware (e.g., different numbers or types of GPUs, varying interconnects) simplifies deployment across diverse environments. Infra teams don't need to homogenize all compute resources, allowing for more flexible resource allocation and utilization.

Tradeoffs and Limitations

While Photon offers significant advantages, it's important to acknowledge its current tradeoffs and limitations:

• Client Memory Requirement: A current constraint is that clients need to be able to hold the full model in memory to participate in the training. This restricts participation to clients with sufficient GPU memory, excluding very small or resource-limited devices. Future work aims to address this by exploring techniques for partial model loading or specialized architectures like sparse models.
• Configuration Complexity: The choice of machine learning configuration, such as the number of local steps and communication frequency, is currently driven by the underlying infrastructure's bandwidth and is not yet fully automated. This requires some initial manual tuning based on network characteristics.
• Scaling to Ultra-Large Models: While Photon has shown excellent results for 13 billion parameter models, scaling to 100 billion parameters and beyond still presents challenges, particularly regarding the communication overhead of transmitting massive parameter chunks. The team is actively researching new optimization procedures to mitigate this, indicating that current methods might not scale indefinitely without further innovation.

Despite these limitations, Photon represents a robust and highly promising framework that fundamentally shifts the paradigm of LLM pre-training, making it more resilient, efficient, and accessible for a broader range of practitioners and organizations.

Key Takeaways

• Photon is the first practical system for Federated LLM pre-training, demonstrating the feasibility and benefits of training large language models across distributed, heterogeneous nodes over the internet.
• It achieves comparable or superior model quality to centralized training while offering significant advantages in efficiency, including 2X higher throughput for a 13B parameter model and up to 90% GPU utilization.
• The system dramatically enhances robustness and fault tolerance, enabling seamless training continuation even when individual client nodes fail, a stark contrast to the fragility of traditional tightly-coupled distributed training.
• Photon unlocks access to vast, decentralized data sources, allowing organizations to train on sensitive or proprietary data without centralization, addressing the impending data scarcity challenge in LLM development.
• Its communication-efficient design is crucial for scaling LLM training over wide-area networks and heterogeneous compute environments, making large-scale AI development more accessible and cost-effective.
• Photon dynamically adapts to hardware heterogeneity, treating strongly connected GPUs differently from weakly connected ones, and tuning batch sizes per device on the fly to maximize local compute efficiency.

About the Speaker(s)

The work on Photon is a collaborative effort, prominently presented by Lorenzo Sani. He is part of a team that includes Alex Iacob, Zeyu Cao, Royson Lee, and Nicholas D. Lane. This diverse group represents a collaboration between Flower Labs, a company focused on federated learning, the Machine Learning Systems Group at the University of Cambridge, and visiting researchers from Bup and Zhejiang University. Their combined expertise in machine learning systems and federated learning has been instrumental in developing this innovative approach to LLM pre-training.