Congestion-Aware Path Selection for Load Balancing in AI Clusters

Abstract

Fast training of large machine learning models requires distributed training on AI clusters consisting of thousands of GPUs. The efficiency of distributed training crucially depends on the efficiency of the network interconnecting GPUs in the cluster. These networks are commonly built using RDMA following a Clos-like datacenter topology. To efficiently utilize the network bandwidth, load balancing is employed to distribute traffic across multiple redundant paths. While there exists numerous techniques for load-balancing in traditional datacenters, these are often either optimized for TCP traffic or require specialized network hardware, thus limiting their utility in AI clusters. This paper presents the design and evaluation of Hopper, a new load-balancing technique optimized for RDMA traffic in AI clusters. Operating entirely at the host level, Hopper requires no specialized hardware or modifications to network switches. It continuously monitors the current path for congestion and dynamically switches traffic to a less congested path when congestion is detected. Furthermore, it incorporates a lightweight mechanism to identify alternative paths and carefully controls the timing of path switching to prevent excessive out-of-order packets. We evaluated Hopper using ns-3 simulations and a testbed implementation. Our evaluations show that Hopper reduces the average and 99-percentile tail flow completion time by up to 20% and 14%, respectively, compared to state-of-the-art host-based load balancing techniques.