【RFC】【大模型训练】原生并行支持动静统一

# 目的 随着AI领域相关技术的发展，神经网络模型参数量及训练数据集的规模均以指数增长。以当下几个热门应用领域为例：在NLP网络中，模型的趋势是通过堆叠transformer以获取更高的精度，如GPT3模型参数达到了千亿级别，以致所需内存远超单卡上限；在人脸识别领域，卷积层的输出经过fc层后实现特征聚合，而fc层的参数大小与图片类别数成正比，对于超大人脸识别网络需处理百万以上规模，必然会受到fc层的内存制约；在推荐网络中，特征头部embedding层可达到百亿特征，也面临同样问题。目前业界主流加速卡的内存只有几十GB，而训练大模型所需内存远超于此，必须把模型切分到多个设备才可运行，因此分布式并行能力是大规模神经网络训练的关键。 综合来看，大模型训练的关键问题主要有：1）受限于单卡内存，如何把大模型切分到多张卡上；2）模型切分到多张卡会引入通信，如何最大化计算通信比；3）模型切分在各个节点中如何调度部署，使得通信模式匹配模型拓扑，发挥带宽。4）在数千节点的集群上，如何提升算法工程师分布式调试调优的效率，另外还要如何考虑降低工程师对大模型进行并行切分的难度。 为解决以上关键问题，MindSpore计划构建一套动静统一的分布式并行技术，目标如下： · 保持单卡逻辑编程，非侵入式修改单卡算法实现分布式并行，提升算法可读性； · 支持pynative执行，通过图编译接口转成整图执行； · 提供三层编程范式，手动、dtensor、shard； · 支持自定义并行与框架原生并行混合编程； · 支持控制流及动态shape； · 支持MPMD图切分及调度； 整体架构图如下：  # 范围 动静统一的分布式并行技术，范围主要包括： - 基础并行： 1）算子并行； 2）自定义并行； 3）优化器并行； 4）pipeline并行； - MPMD： 1）异构图切分/调度； 2）卡内多核并行； - 重计算/swap； # 关键特性设计 ## 张量排布表达 ### 基本原理 我们使用Layout来表达张量的分布式排布，内部包含3个成员： · device_matrix: 描述集群中卡号的逻辑排布方式 · alias_name：device_matrix中每个轴的别名 · rank_list（可选）：device_matrix对应的设备列表 比如： layout = Layout((4, 2), ("dp", "mp"), (0，1，2，3，4, 5, 6, 7)) 设备按rank号的逻辑排布如下：  对于一个二维张量： 可以通过layout来表达它的各种切分形态：  - 注：各种颜色的内容代表不同的卡上所拥有的数据，比如layout("dp", "mp")中橙色的内容，代表rank0上拥有的数据内容为0和2。 各种切分形态之间可以通过重排进行转换，示意如下：  ### 接口设计 layout = Layout(device_matrix, alias_name, rank_list=None) - 输入： · device_matrix: 描述集群中卡号的逻辑排布方式，如果rank_list为None，则卡号按0开始 · alias_name：device_matrix中每个轴的别名 · rank_list（可选）：device_matrix对应的设备列表 - 输出： 返回一个Layout对象实例 通过对Layout对象实例传入轴别名列表，能指定对应的切分表达，如： layout = Layout((2, 2), ("dp", "mp")) layout("dp", "None")表示将张量的行切分到设备矩阵的“dp”轴， 而列不切分。 注：“None”是特殊的轴别名，使用时表示该维度不切分 ### 需求拆解 - tensor增加layout结构 - local tensor与global tensor转换 - 重排布 相关pr链接： h…