Accelerating Heterogeneous Tensor Parallelism via Flexible Workload Control

Abstract

Transformer-based foundation models are becoming deeper and larger. For fast training, their billions of parameters (tensors) are split onto parallel tasks running on many modern yet expensive accelerators. To amortize the huge hardware investments, it is cost-effective to share the aggregated resources among multi-tenants. However, resource contention yields the heavy straggling problem. Existing works feature contributions for the traditional data parallelism. They cannot work well for the new tensor parallelism, due to the dependency among split tensors. This paper is the first attempt on accelerating heterogeneous tensor parallelism. We summarize specific challenges, including the very frequent synchronizations and the heavy tensor computation workloads. Our solution is to resize dimensions of parameters on demand, to quickly and dynamically balance workloads. The accuracy loss is reduced through priority resizing. We also migrate workloads between tasks, without any loss of accuracy. The most efficient communication primitives are selected and then scheduled in a non-redundant manner, to reduce the runtime latency. Our final hybrid solution is built on top of resizing and migration. By studying the tradeoff between accuracy and efficiency, it can smartly hit the “sweet spot”. Extensive experiments validate the effectiveness of our proposals.