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.
翻译:深度学习正经历基础模型的兴起,这些模型有望引领各个领域。海量参数使得基础模型训练必须采用张量模型并行(TMP)。然而,TMP需要频繁的通信操作,这显著降低了训练效率。本文提出Oases——一种具有重叠通信的自动化TMP方法,用于加速基础模型训练。Oases提出了一种细粒度训练调度机制,以最大化重叠通信与具有数据依赖性的计算操作。此外,我们设计了Oases规划器,用于搜索最优模型并行策略以实现进一步加速。与现有方法不同,Oases规划器专门针对重叠通信-计算操作的成本建模。我们在多种模型设置和训练环境下评估Oases,并将其与四种最先进的实现进行对比。实验结果表明,Oases相较于最快基线实现了1.01-1.48倍的加速比,相较于Megatron-LM的加速比高达1.9倍。