Accelerating Distributed Training on Parameter Server Architecture With Path-Aware Multicast

It is observed that the bottleneck in distributed training has shifted from computation to communication due to contention in concurrent transmissions and substantial redundant traffic. In the Parameter Server (PS) architecture, the server aggregates gradients from multiple workers and then distributes updated model parameters back to the workers in a one-to-many manner. Currently, model parameters are distributed via unicast, sending multiple identical copies of the data, which leads to significant bandwidth waste. Although multicast can save bandwidth, current approaches have two main drawbacks: on one hand, many protocols require maintaining excessive multicast state inside the network; on the other hand, the lack of coordination among multiple multicast trees can still lead to path conflicts. In this work, we propose path-aware multicast, which includes innetwork multicast tree reservation and per-hop control multicast. Specifically, before each round of model parameter distribution, the server queries the network for a multicast tree that satisfies the bandwidth requirement. The calculated multicast tree is then returned with bandwidth reserved at its tree nodes. Next, model parameters are forwarded with hop-by-hop control along the multicast tree. After the multicast is completed, the reserved network resources are released. Our evaluation shows that in an <tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$8 \times 8$</tex> spine-leaf topology, path-aware multicast improves link load balancing by 32.6 % compared to random multicast and accelerates model parameter distribution by up to nearly <tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$N \times$</tex> compared to unicast, where <tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$N$</tex> is the number of workers.