arXiv preprint ·
Accelerating Sharded Data Parallelism at Scale with Federated Learning
- Published
- Source
- arXiv preprint, posted before peer review
- DOI
- 10.1007/978-3-032-35251-4_30
- Original
- Read it at arXiv
- Topics
Abstract
The symbiotic scaling of artificial intelligence models and high-performance computing systems continually creates algorithmic challenges in their convergence. Foundation models (FMs) are a crucial example, requiring months-long training on thousands of cutting-edge GPUs. Sharded data parallelism (DP) is the dominant strategy to accelerate such computations by splitting data and models across multiple GPUs. However, it incurs prohibitive communication overhead when deployed at scale, particularly on multi-tier interconnects with heterogeneous performance. Inspired by the efficient communication principles of federated learning (FL), this work introduces two hybrid algorithms - FL+FSDP and FL+HSDP - interleaving sharded DP with FedAvg-style aggregations. Such approaches decouple large DP deployments into smaller, loosely-coupled federation groups, requiring minimal inter-group traffic while keeping the global batch size bounded by the groups' size. Formal analysis of communication costs and experimental validation prove their scalability and flexibility. A Llama3.1 8B pre-training on 512 A100 GPUs shows that, under identical hyperparameters, FL+FSDP and FL+HSDP achieve up to 8.04 faster data processing and 4.48 lower evaluation perplexity than their counterparts, demonstrating superior computational efficiency and improved model quality. These properties stem from reduced communication overhead and the bounded growth of the global batch size relative to the federation group size.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Distributed & Parallel Computing
PrefixBench-H100: Characterizing Prefix Reuse and Time-to-First-Token in H100 LLM Serving
Hybrid GPU-CPU Retrieval for Personalized Search at Ultra-Large Scale
Understanding and Exploiting Diagonal Attention Sparsity in Autoregressive Image Generation
Cross-Platform vs Native Mobile Development: An Empirical Study of Software Quality Trade-offs
The Life of a Token: from Words to Bits on the Wire
Get the few papers worth your time in Distributed & Parallel Computing every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.