  Distributed training is a solution to reduce DNN training time by splitting
the task across multiple NPUs (e.g., GPU/TPU). However, distributed training
adds communication overhead between the NPUs in order to synchronize the
gradients and/or activation, depending on the parallelization strategy. In
next-generation platforms for training at scale, NPUs will be connected through
multi-dimensional networks with diverse, heterogeneous bandwidths. This work
identifies a looming challenge of keeping all network dimensions busy and
maximizing the network BW within the hybrid environment if we leverage
scheduling techniques for collective communication on systems today. We propose
Themis, a novel collective scheduling scheme that dynamically schedules
collectives (divided into chunks) to balance the communication loads across all
dimensions, further improving the network BW utilization. Our results show that
on average, Themis can improve the network BW utilization of the single
All-Reduce by 1.72X (2.70X max), and improve the end-to-end training iteration
performance of real workloads such as ResNet-152, GNMT, DLRM, and
Transformer-1T by 1.49X (2.25X max), 1.30X (1.78X max), 1.30X (1.77X max), and
1.25X (1.53X max), respectively.
