TaLB: Tensor-aware Load Balancing for Distributed DNN Training Acceleration

Increasingly large-scale models and rich data sets make communication overhead a key bottleneck for distributed Deep Neural Network (DNN) training, constantly attracting the attention of academia and industry. Despite continuous efforts, prior solutions such as pipelining computation/communication and in-network gradient compression/scheduling do not focus on how to accelerate DNN training through load balancing in datacenter networks (DCNs). However, the existing load balancing mechanisms are unaware of tensor integrity and priority for gradient parameter synchronization during the DNN training iterations, resulting in severe tensor tail latency and slow model convergence speed. In this paper, we present a Tensor-aware Load Balancing (TaLB) scheme to accelerate DNN training. Specifically, TaLB identifies the different priority tensors and makes (re)routing decisions based on the tensor-level granularity to cut the high-priority tensors tail delay. The testbed implementation and large-scale NS-3 simulation results show that TaLB effectively accelerates DNN training speed. For example, TaLB significantly reduces the average flow completion time (FCT) by up to 55%, and accelerates the model training speed up to 2.37× on VGG19, ResNet50 and AlexNet models.