【AI系统】昇思MindSpore并行

昇思MindSpore并行

本文将会介绍昇思MindSpore的并行训练技术，以及如何通过张量重排布和自动微分简化并行策略搜索，实现高效大模型训练。

大模型的带来

随着深度学习的发展，为了实现更高的准确率和更丰富的应用场景，训练数据集和神经网络模型的规模日益增大。特别是自然语言处理（Natural Language Processing，NLP）领域，数据集的范围从200MB到541TB不等。模型的尺寸也从BERT的3.4亿个参数，Transformer-xl的8亿个参数，GPT-2的150亿个参数，到SwitchTransformer的万亿参数，单独存储万亿模型的参数则需要占用TB级别空间。

主流的用于训练的加速器（GPU、TPU和Ascend）的内存仍然只有几十GB，因此，为了能够且高效地完成大模型训练，深度学习框架需要提供丰富的并行训练的能力，支持数据并行、模型并行、混合并行、流水线并行、异构训练和专家并行等技术。

AI 框架的局限性

当前的主流框架（如TensorFlow、Caffe和MXNet）提供了一些基本的并行技术，其中大多数框架提供了算子级别切分、流水线并行或者优化器切分的功能，支持的并行维度和功能的完整度欠佳。

• 第一，这些框架通过手动的切分来切分神经网络模型来实现模型并行，配置难度非常大，对开发者的要求非常高，需要有丰富经验的专家来操作。实现混合并行（数据并行和模型并行同时进行）又极大增加了开发的复杂度。最近的研究成果提出了简化混合并行的方法，但这些方法在几个方面都存在局限性。

• 第二，随着目前模型规模拓展到了万亿级别，训练卡数的规模也上升到了千级，以往的算子级别切分的并行不能满足目前大模型的需求。在大集群训练中，由于模型切分导致的通信占比在整个迭代耗时中升高，需要引入流水线并行等计算来降低通信占比，提升模型训练效率。另外，混合专家（MoE）技术能够在提升模型规模的同时，较少的提升计算量，是目前的一种主流技术。

MindSpore 大模型并行

昇思MindSpore的目标是提供完整的大模型训练落地流程。为了解决前后设备张量排布不一致的问题，在并行化策略搜索中引入了张量重排布（Tensor Redistribution，TR），这使输出张量的设备布局在输入到后续算子之前能够被转换。

但是，对于复杂大型模型的搜索并行策略，在考虑张量重排布时，需要克服的挑战主要有两个。首先，既然张量重排布将通信算子（例如AllGather）引入数据流图，那么如何像普通算子一样自动地对通信算子求导呢？对于每个相应的前向算子，都需要获取反向算子，用来更新可训练参数。目前的框架需要专家在反向阶段手动添加SEND和RECV源语来传递梯度，这对模型开发者来说是一项具有挑战性的工作，尤其是在模型比较复杂的情况下。

其次，随着张量重排布对策略空间的极大扩展，如何为复杂的大型模型高效地找到一个好的策略？在功能和效率方面，该算法需要为具有非线性结构的大模型快速找到一种策略。在性能方面，算法返回的策略应该会缩短端到端的训练时间。需要对运行成本进行仔细建模，这一过程也增加了人工成本。最后，如何解决万亿模型训练的实际问题，最大化的利用硬件算力，同时降低训练成本？

MindSpore 并行策略

昇思MindSpore针对上述三个挑战，推出了完整的大模型训