Topology-Aware All-Reduce Algorithm Synthesizer Based on Deep Reinforcement Learning

With the rapid growth of the scale of distributed machine learning (ML) systems, All-Reduce, enabling the exchange and reduction of chunks among a group of nodes, has emerged as a critical bottleneck limiting the training efficiency of ML models. Although various collective communication libraries (CCLs) have been developed based on predesigned All-Reduce algorithms, they have been proven inefficient in distributed clusters with asymmetric and heterogeneous topologies. To address this challenge, this paper utilizes deep reinforcement learning (DRL) to develop a topology-aware All-Reduce algorithm synthesizer, namely, RLTA. It transforms the All-Reduce operation into a minimum-time network traversal problem of chunks. Through customizing the state, action, and reward function within the DRL framework, the agent of RLTA autonomously learns a multi-chunk-path plan for chunks, which can be directly executed as an enhanced All-Reduce algorithm to improve the training efficiency of distributed computing clusters. Finally, we have justified the effectiveness of RLTA across diverse network topologies and chunk sizes, and our proposed scheme can achieve an average 2.5× communication speedup and 92.2% link utilization over benchmarking All-Reduce algorithms.