Tesseract: Parallelize the Tensor Parallelism Efficiently

TL;DR

Tesseract introduces a novel tensor parallelism method that significantly improves efficiency and scalability for training large deep learning models on limited GPU resources by reducing communication overhead and increasing memory capacity.

Contribution

The paper proposes Tesseract, a new tensor parallelism design that enhances scalability and efficiency by adding a novel dimension, reducing communication costs, and increasing memory capacity compared to previous methods.

Findings

Achieves 1.38x and 1.53x speedups over 1-D and 2-D methods.

Realizes up to 4.0x inference speedup and 3.4x throughput improvement.

Reduces communication overhead, enabling training of larger models with limited GPU resources.

Abstract

Together with the improvements in state-of-the-art accuracies of various tasks, deep learning models are getting significantly larger. However, it is extremely difficult to implement these large models because limited GPU memory makes it impossible to fit large models into a single GPU or even a GPU server. Besides, it is highly necessary to reduce the training time for large models. Previous methods like Megatron-LM implemented a 1-Dimensional distributed method to use GPUs to speed up the training. However, these methods have a high communication overhead and a low scaling efficiency on large-scale clusters. To solve these problems, we propose Tesseract, a highly scalable tensor parallelism with a novel design. It increases efficiency by reducing communication overhead and lowers the memory required for each GPU. By introducing the novel dimension into tensor parallelism, Tesseract greatly increases the memory capacity of tensor parallelism. Concretely, this new dimension furthermore increases the degree of tensor parallelism. Compared to previous 1-D and 2-D methods, Tesseract manages to reduce the communication cost on each layer, resulting in speedups of 1.38x and 1.53x respectively with strong scaling. In weak scaling experiments, Tesseract achieves a maximum of 4.0/1.7 times inference speedup and 3.4/1.7 times throughput improvement compared to 1-D/2-D methods, respectively. By introducing Tesseract, we offer a more efficient and scalable way to implement large deep learning models with limited GPU resources.