Photon: Federated LLM Pre-Training (2411.02908v1)

Abstract: Scaling 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.

• The paper introduces a federated framework that pre-trains LLMs up to 7B parameters with 20% higher throughput and 64×–512× lower communication overhead compared to centralized approaches.
• The paper demonstrates robust optimization using federated averaging, which enables twice the convergence speed through high learning rates and small batch sizes despite data heterogeneity.
• The paper showcases architectural innovations that efficiently utilize decentralized compute resources and pave the way for personalized models while reducing reliance on centralized data centers.

Overview of "Photon: Federated LLM Pre-Training"

This paper introduces Photon, an open-source system designed for federated pre-training of LLMs across low-bandwidth distributed systems. Traditional centralized training of LLMs is resource-intensive, often constrained by the necessity for high-bandwidth connections within large data centers. In contrast, Photon's federated framework allows for the collaborative training of models across decentralized devices, potentially overcoming the limitations of centralized infrastructure. The authors claim that Photon can pre-train models up to 7 billion parameters, achieving lower perplexity scores than centralized counterparts and demonstrating significant reductions in communication overhead.

Photon leverages federated learning (FL) techniques, which enable training on disparate data sources without the data leaving its original location. This could democratize access to training large-scale models by utilizing globally distributed computational resources, including personal hardware accelerators. Photon adapts to varying connectivity scenarios using techniques like federated averaging, which provides robustness to hyperparameter variations and leverages small batch sizes at high learning rates to achieve twice the convergence speed of methods like DiLoCo.

Key Contributions and Results

1. Scalability and Performance: Photon successfully pre-trains LLMs up to 7 billion parameters across distributed nodes. The federated models reach better perplexity scores compared to those trained in traditional centralized settings. The system achieves up to 20% higher throughput with significantly reduced communication requirements (64×–512× less) by optimizing for low-bandwidth connections.
2. Robust Optimization: By utilizing federated averaging, Photon maintains robust optimization paths despite data heterogeneity and variations in local hardware capabilities. The design allows models to converge rapidly, which the authors attribute to a combination of small batch sizes and high learning rates.
3. Efficient Use of Compute Resources: As more computational resources are augmented to the federated system, Photon effectively scales and reduces training time, providing a compute-time trade-off similar to centralized methods. This is especially significant given the constraints of federated systems concerning communication and synchronization.
4. Architectural Innovations: Photon integrates adaptive local parallelism, allowing a seamless adaptation between traditional distributed training methodologies and low-bandwidth federated learning approaches. This ensures optimal resource utilization tailored to the specific capabilities and connectivity of each participating node.

Implications and Future Directions

Photon presents a promising approach to training LLMs in a decentralized manner, reducing the dependency on centralized data centers and potentially lowering the barrier for entry in training large models. The significant reduction in communication required and its robustness to heterogeneity opens new opportunities for using diverse datasets from different geographies without infringing on data privacy.

The implications of Photon's approach are multifold. Federated LLM training could spur the development of personalized models that adapt based on localized datasets, improving their capability to process language in contextually relevant ways. Furthermore, this method mitigates the challenges of exorbitant computational and environmental costs associated with scaling data centers for LLM training.

For future developments, exploring enhanced personalization, continual learning, and addressing the nuances of data heterogeneity remain critical. Investigating adaptive mechanisms for federated hyperparameter tuning could enhance model robustness and generalizability further. Another avenue is deploying Photon in cross-device federated settings, necessitating optimizations like parameter-efficient fine-tuning and advanced quantization techniques to tackle the challenges posed by limited computational capabilities prevalent in mobile and IoT devices.

Overall, Photon represents a significant advance in the landscape of distributed LLM training, offering a scalable and efficient solution adaptable to the constraints and opportunities inherent in a federated paradigm.