STM-WFBP: Selective Tensor Merging Method in Distributed Learning
摘要
Communication is an important constraint for distributed deep learning training, and efficient communication scheduling methods such as tensor merging accelerate distributed training by merging tensors from model layers. However, existing optimizations perform tensor merging greedily, failing to consider the occurrence of tensor stagnation for successive merging of multiple layers, resulting in suboptimal training performance. In this paper, we propose the selective tensor merging method STM-WFBP to guide the tensor merging process. This method intelligently determines whether the solvable layers are effectively merged and selects the effectively merged layers, thus avoiding the phenomenon of tensor retention. Experimental results on two GPU clusters show that the STM-WFBP achieves a 4%-12% optimization in iteration time over existing methods, which proves its effectiveness.