🤖 AI Summary
To address the high communication overhead and limited compute-communication overlap in tensor model parallelism (TMP) for large language models, this paper proposes Oases, a dependency-aware, fine-grained TMP auto-optimization framework. Our method introduces: (1) operator-level fine-grained training scheduling; (2) communication-computation overlap–aware modeling and planning; (3) the first automated TMP strategy search algorithm supporting dependency constraints; and (4) end-to-end runtime optimization for TMP. Evaluated across multiple models (e.g., LLaMA, BERT) and hardware platforms (A100/H100), Oases achieves 1.01–1.48× speedup over state-of-the-art approaches and up to 1.9× improvement over Megatron-LM. These results demonstrate significant gains in training efficiency for large-scale foundation models.
📝 Abstract
Deep learning is experiencing a rise in foundation models that are expected to lead in various fields. The massive number of parameters necessitates the use of tensor model parallelism (TMP) in foundation model training. However, TMP requires frequent communication operations which significantly reduces the training efficiency. In this paper, we present Oases, an automated TMP method with overlapped communication to accelerate foundation model training. Oases proposes a fine-grained training schedule to maximize overlapping communication and computation operations that have data dependence. Additionally, we design the Oases planner that searches for the best model parallel strategy to achieve further accelerations. Unlike existing methods, Oases planner is specifically tailored to model the cost of overlapped communication-computation operations. We evaluate Oases on various model settings and train environments, and compare Oases to four stat-of-the-art implementations. Experimental results demonstrate that Oases achieves speedups of 1.01--1.48X over the fastest baseline, and speedups of up to 1.9X over Megatron-LM.
