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.
- Tags
- communication-efficiencydistributed-trainingfederated-learningfsdpgpu-scalingllamasharded-data-parallelism
- Vendors
- Meta
- Products
- Llama 3.1 8B
- AI models
- Llama 3.1 8B
From the original article
1 min read at arxiv.orgThe 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.
By Gianluca Mittone · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.