DSA: Efficient Data-Plane Memory Scheduler for In-Network Aggregation to Accelerate Distributed Training

To reduce the traffic volume and accelerate communication in distributed training (DT) jobs, recent works introduce In-Network Aggregation (INA) to move the gradient summation into network programmable switches. However, switch memory is a scarce resource, unable to support massive DT jobs in data centers, and existing INA solutions have not utilized switch memory to the best extent. We propose DSA, an Efficient Data-Plane switch memory Scheduler for in-network Aggregation. DSA introduces preemption to the switch memory management for INA jobs. Furthermore, under packet preemption scenarios, DSA optimizes the selective retransmission mechanism to reduce redundant retransimtting packets to alleviate congestion. In the data plane, DSA allows gradient tensors with high priority to preempt the switch aggregators (basic computation unit in INA) from tensors with low priority, which avoids an aggregator wasting time in idle. In the control plane, DSA devises a priority policy which assigns high priority to gradient tensors that benefit overall job efficiency more, e.g., communication-intensive jobs. We implement the prototype of DSA. The experimental results show that DSA can improve the average job completion time (JCT) by up to $1.35\times $ compared with baseline solutions.