一种混合张量-专家-数据并行方法以优化混合专家训练
A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training
AI总结:
提出DeepSpeed-TED三维混合并行算法,结合数据、张量和专家并行,支持训练比现有技术大4-8倍基础模型的MoE模型,并在400亿参数模型上实现26%加速。
摘要:
混合专家(MoE)是一种神经网络架构,它在基础模型上添加了稀疏激活的专家模块,从而在不增加计算成本的情况下增加参数数量。然而,当前的分布式深度学习框架在训练具有大型基础模型的高质量MoE模型方面能力有限。在这项工作中,我们提出了DeepSpeed-TED,一种新颖的三维混合并行算法,它结合了数据、张量和专家并行,使得能够训练比当前最先进技术大4到8倍基础模型的MoE模型。我们还描述了优化器步骤中的内存优化,以及消除不必要数据移动的通信优化。我们在DeepSpeed中实现了我们的方法,并在128块V100 GPU上训练一个400亿参数的MoE模型(67亿基础模型,16个专家)时,相比基线(即没有我们的通信优化)实现了26%的加速。
英文摘要:
Mixture-of-Experts (MoE) is a neural network architecture that adds sparsely activated expert blocks to a base model, increasing the number of parameters without impacting computational costs. However, current distributed deep learning frameworks are limited in their ability to train high-quality MoE models with large base models. In this work, we present DeepSpeed-TED, a novel, three-dimensional, hybrid parallel algorithm that combines data, tensor, and expert parallelism to enable the training of MoE models with 4 to 8x larger base models than the current state-of-the-art. We also describe memory optimizations in the optimizer step, and communication optimizations that eliminate unnecessary data movement. We implement our approach in DeepSpeed and achieve speedups of 26% over a baseline (i.e. without our communication optimizations) when training a 40 billion parameter MoE model (6.7 billion base model with 16 experts) on 128 V100 GPUs.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。