Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Models

TL;DR

Themis is a novel collective scheduling policy that dynamically balances communication loads across heterogeneous network dimensions, significantly improving bandwidth utilization and training performance in distributed deep learning systems.

Contribution

Themis introduces a dynamic collective scheduling scheme that optimizes network bandwidth utilization across multiple dimensions in heterogeneous environments.

Findings

Improves network bandwidth utilization by 1.72X on average.

Enhances training iteration performance for various workloads by up to 2.25X.

Effectively balances communication loads across network dimensions.

Abstract

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.