6 papers
Decoupled DiLoCo for Resilient Distributed Pre-training
Arthur Douillard, Keith Rush, Yani Donchev +14
Modern large-scale language model pre-training relies heavily on the single program multiple data (SPMD) paradigm, which requires tight coupling across accelerators. Due to this co…
Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities
Gheorghe Comanici, Eric Bieber, Mike Schaekermann +3431
In this report, we introduce the Gemini 2.X model family: Gemini 2.5 Pro and Gemini 2.5 Flash, as well as our earlier Gemini 2.0 Flash and Flash-Lite models. Gemini 2.5 Pro is our…
Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo
Zachary Charles, Gabriel Teston, Lucio Dery +5
As we scale to more massive machine learning models, the frequent synchronization demands inherent in data-parallel approaches create significant slowdowns, posing a critical chall…
Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch
Arthur Douillard, Yanislav Donchev, Keith Rush +11
Training of large language models (LLMs) is typically distributed across a large number of accelerators to reduce training time. Since internal states and parameter gradients need…
Federated Automatic Differentiation
Keith Rush, Zachary Charles, Zachary Garrett
Federated learning (FL) is a general framework for learning across an axis of group partitioned data (heterogeneous clients) while preserving data privacy, under the orchestration…
DrJAX: Scalable and Differentiable MapReduce Primitives in JAX
Keith Rush, Zachary Charles, Zachary Garrett +2
We present DrJAX, a JAX-based library designed to support large-scale distributed and parallel machine learning algorithms that use MapReduce-style operations. DrJAX leverages JAX'…