Accelerating Sharded Data Parallelism at Scale with Federated Learning
Hybrid FL+FSDP and FL+HSDP algorithms cut communication overhead in sharded data parallelism, accelerating Llama3.1 8B pre-training on 512 A100 GPUs by up to 8x.
The paper introduces FL+FSDP and FL+HSDP, hybrid algorithms that interleave sharded data parallelism with FedAvg-style federated aggregations to decouple large DP deployments into loosely-coupled federation groups. Formal communication-cost analysis and experiments on multi-tier GPU interconnects demonstrate scalability and flexibility. A Llama3.1 8B pre-training run on 512 A100 GPUs achieves up to 8.04x faster data processing and 4.48 lower evaluation perplexity than sharded DP baselines under identical hyperparameters.
- Introduces FL+FSDP and FL+HSDP, interleaving sharded data parallelism with FedAvg-style aggregations across loosely-coupled federation groups.
- Llama3.1 8B pre-training on 512 A100 GPUs achieves up to 8.04x faster data processing than DP counterparts.
- Shows 4.48 lower evaluation perplexity under identical hyperparameters, with formally analyzed communication costs.
- Targets prohibitive communication overhead of sharded DP on heterogeneous multi-tier interconnects.
Full article197 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.20359