Accelerating Communication for Distributed Training with Collaborative In-network Aggregation

Distributed training (DT) has emerged as a crucial and popular practice for training high-quality machine learning (ML) models. However, the communication efficiency of gradient aggregation during the training process has become the primary performance bottleneck. Advanced programmable switches with in-network computing capabilities offer a promising solution to improve the communication efficiency of DT by offloading partial gradient aggregation tasks from servers to switches within the network. In this paper, we propose CINA (Collaborative In-network Aggregation), a novel approach that harnesses the in-network computing capabilities of programmable switches and allows multiple switches and servers to collaboratively perform gradient aggregation tasks, thereby improving the communication efficiency of distributed machine learning training. CINA aims to minimize the aggregation completion time by jointly optimizing the aggregation task assignment among switches and the gradient routing of aggregation requests. We formulate this joint optimization problem as an integer nonlinear programming problem under multi-dimensional resource constraints and design an efficient two-stage random rounding algorithm to compute satisfactory solutions rapidly. Extensive experimental results demonstrate that CINA significantly outperforms other state-of-the-art in-network gradient aggregation solutions, reducing the aggregation completion time by up to 75%.