Accelerating Sharded Data Parallelism at Scale with Federated Learning

As foundation models (FMs) grow in size, training them requires thousands of GPUs, which often leads to prohibitive communication overheads in standard sharded data parallelism (DP) setups. This paper investigates whether incorporating federated learning (FL) principles into sharded DP can mitigate these bottlenecks by reducing the frequency of heavy collective communications across slow, noisy inter-node links.

The authors introduce two hybrid algorithms, FL+FSDP and FL+HSDP, which interleave traditional sharded DP with FedAvg-style aggregations. By partitioning the global process pool into smaller federation groups, the algorithms confine expensive collective operations to intra-group communications, while using lightweight, periodic aggregations for inter-group synchronization. The authors provide a formal communication cost analysis and validate the approach by pre-training a Llama3.1 8B model on 512 A100 GPUs across 128 nodes.

The proposed FL-augmented methods demonstrate superior computational efficiency and model quality compared to standard FSDP and HSDP. Specifically, FL+FSDP achieved up to 8.04x faster data processing, while FL+HSDP yielded a 4.48x lower evaluation perplexity. These improvements occur because the FL-based approach bounds the growth of the global batch size relative to the federation group size, leading to more stable convergence, and restricts high-latency collective operations to faster, local interconnects.

This work establishes federated learning as a viable, high-level primitive for intra-cluster distributed training, not just for geographically dispersed devices. By providing a scalable communication model that reflects the hierarchical nature of modern HPC interconnects, these algorithms offer a path to train increasingly large foundation models without being limited by the performance of the slowest network links.