In-Network AllReduce Optimization with Virtual Aggregation Trees
Haoyu Song · 2024
AllReduce is a critical performance bottleneck for distributed deep learning and large model training in data centers for AI computing. In-Network Aggregation (INA) has been identified as an effective accelerating technique to improve its performance. However, the existing schemes either use only the ToR switches for aggregation or rely on physical aggregation trees, which limit the efficiency and scalablility of INA. In this paper we propose a novel method to construct Virtual Aggregation Trees (VAT) on top of the physical data center network topology to take advantage of all network switches and balance their load. We show that VAT encoding can be realized by repurposing the IETF BIER bitmap. Thus, by integrating with the BIER multicast feature for aggregation result distribution, we have a complete AllReduce solution with network-facilitated collective communication optimizations. We develop a greedy algorithm for incremental VAT construction. The evaluations show that VAT outperforms the other state-of-art INA schemes in terms of the number of parallel jobs supported on networks with the Spine-Leaf or Fat-tree topologies.