Ofan improves AI training speed by balancing network load better
Ofan: Optimal Load Balancing for AI Training
Networking and Internet ArchitectureDistributed, Parallel, and Cluster ComputingPerformance
Summary
Training large AI models requires fast and balanced data movement across network switches. Existing methods can cause traffic jams that slow down training. The authors show that current algorithms have fundamental limits in handling such traffic efficiently on fat tree networks. They introduce Ofan, a new way for network switches to be smart about where data goes, greatly reducing delays. Testing with large AI models shows Ofan can make training run much faster compared to previous methods.
What this means in practice
- •For data center network engineers: Optimize network switch routing to reduce delays during large AI model training, improving overall throughput and efficiency.
- •For cloud infrastructure operators: Deploy Ofan to better balance heavy network loads in AI workloads, lowering training times for large-scale machine learning services.
Authors
Sarah McClure, Evyatar Cohen, Jakob Krebs, Alexander Shpiner, Mark Silberstein, Scott Shenker, Sylvia Ratnasamy, Isaac Keslassy
Abstract
The extreme collective completion time (CCT) demands of AI workloads challenge existing packet spraying algorithms, which can have trouble efficiently load-balancing workloads that are sent at full line rates. We trace this to a structural cause: on a fat tree, once a packet picks its upward path, the downward path to its destination is unique, so destination-oblivious schemes cannot undo the imbalance it creates. We prove that such schemes can suffer from $Θ(\sqrt{m})$ queueing for messages of size $m$, thus eventually triggering rate reductions by the congestion control. Instead, we suggest Ofan, a switch-based destination-aware LB scheme that can reach $O(1)$ queueing. We also present its pOfan variant that fits the pipe-based architecture of current switches. Our P4 implementation shows that it consumes modest resources. An end-to-end FSDP2 evaluation with Llama-3 405B-parameter models shows that Ofan cuts CCT inflation by $16$--$39\times$ when compared to existing algorithms.