Gianluca Mittone, Marco Aldinucci
4 min
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.
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.
Sam: [nodding] So the cost you're paying isn't latency anymore, it's tuning effort. You've moved the engineering problem from the network layer to the training configuration.
Alex: [analytical] That's exactly the right framing. And for anyone running foundation model training on a congested supercomputer fabric — where the interconnect between nodes is genuinely the bottleneck — that's a trade worth making. The alternative is that your GPUs spend a substantial fraction of their time waiting on slow cross-node all-reduces. This approach lets you keep the GPUs busy doing actual computation, and you pay for it with periodic, coarser synchronization. [[RP_SECTION:hardware-topology-and-scaling|Hardware Topology and Scaling]]
Sam: [probing further] One thing I'd want to know as a referee: how sensitive are these gains to the specific hardware topology? If your intra-node bandwidth advantage over inter-node isn't that large, does the whole argument collapse?
Alex: [measured] That's the constraint that most limits generalizability. The speedup numbers are impressive on the hardware they tested, but the architecture is explicitly designed to exploit a bandwidth differential between intra- and inter-node links. On a system with more uniform interconnects — say, a tightly integrated GPU cluster with high-speed fabric across all nodes — the benefit of confining collectives to the group level shrinks. The paper doesn't fully characterize where that crossover point is, which is the gap a careful follow-up would need to close. [[RP_SECTION:practical-implementation-takeaways|Practical Implementation Takeaways]]
Sam: [reflecting] So the practical takeaway is fairly specific: if you're training on infrastructure where cross-node bandwidth is genuinely scarce relative to intra-node, this hybrid FL-plus-FSDP topology gives you a concrete path to better scaling. But you need to profile your federation frequency, and you need to know your hardware's bandwidth profile before you can predict how much of that speedup you'll actually recover.
Alex: [concluding, steady] That's it. The contribution is reframing a hardware limitation as a configurable parameter. Whether the tuning overhead is worth it depends on your setup — but for researchers stuck on congested fabrics who've already exhausted the standard FSDP playbook, this gives you a principled alternative rather than just throwing more hardware at the problem. Thanks for listening to ResearchPod.