昇思MindSpore大模型并行训练全解析：张量重排布、通信算子自动微分与多级并行策略

文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载本文系统梳理昇思MindSpore在大模型并行训练领域的完整解决方案从大模型时代面临的参数规模与单卡内存矛盾出发讲解张量重排布Tensor RedistributionTR如何自动解决设备间张量布局不一致问题并深入剖析算子级别并行、函数式算子切分、流水线并行、MoE专家并行与异构并行训练等关键技术帮助开发者理解自动并行的底层原理掌握通过shard配置、jit装饰符等接口落地大模型训练实战的能力。大模型带来的训练挑战随着深度学习的发展为了追求更高的准确率和更丰富的应用场景训练数据集和神经网络模型的规模日益增大。以自然语言处理Natural Language ProcessingNLP领域为例训练数据集的规模从 200MB 到 541TB 不等模型参数规模也从 BERT 的 3.4 亿个、Transformer-XL 的 8 亿个、GPT-2 的 150 亿个一路扩展到 Switch-Transformer 的万亿级别——单独存储一个万亿参数的模型就需要 TB 级别的空间。而主流的训练加速器GPU、TPU 和昇腾 Ascend的显存/内存仍然只有几十 GB。单个设备既装不下模型参数也扛不住超大训练集的计算量因此深度学习框架必须提供丰富的并行训练能力支持数据并行、模型并行、混合并行、流水线并行、异构训练和专家并行MoE等技术才能高效完成大模型训练。AI 框架的局限性当前的主流框架如 TensorFlow、Caffe 和 MXNet虽然提供了一些基本的并行技术其中大多数框架提供了算子级别切分、流水线并行或者优化器切分的功能但支持的并行维度和功能完整度欠佳手动切分配置难度大这些框架通过手动切分神经网络模型来实现模型并行配置难度非常大对开发者要求极高通常需要有丰富经验的专家来操作。而要实现数据并行与模型并行同时进行的混合并行又会极大增加开发复杂度。虽然近年来的研究成果提出了简化混合并行的方法但这些方法在并行维度覆盖、易用性等方面仍然存在局限。算子级切分难以满足万亿级模型随着模型规模拓展到万亿级别训练卡数规模也上升到了千卡级别以往的算子级别切分并行已不能满足当前大模型的需求。在大集群训练中模型切分导致的通信占比在整个迭代耗时中不断升高需要引入流水线并行等手段来降低通信占比、提升训练效率。与此同时混合专家MoE技术能够在提升模型规模的同时仅少量增加计算量已成为主流技术方向之一。MindSpore 大模型并行三大挑战与整体方案昇思MindSpore的目标是提供完整的大模型训练落地流程。针对上述问题MindSpore 给出的核心思路是在并行化策略搜索中引入张量重排布Tensor RedistributionTR使输出张量的设备布局在输入到后续算子之前能够被自动转换从而解决前后设备张量排布不一致的问题。但要在考虑张量重排布的前提下为复杂大型模型搜索并行策略需要克服三个核心挑战通信算子的自动微分张量重排布会将通信算子如 AllGather引入数据流图那么如何像普通算子一样自动地对通信算子求导对每个前向算子都需要获取对应的反向算子用来更新可训练参数。传统框架需要专家在反向阶段手动添加 SEND 和 RECV 原语来传递梯度这对模型开发者是一项极具挑战性的工作尤其在模型比较复杂时。高效策略搜索张量重排布极大地扩展了策略空间如何为复杂的大型模型高效地找到好的策略在功能上算法需要为具有非线性结构的大模型快速找到策略在性能上算法返回的策略应当能缩短端到端训练时间这要求对运行成本进行仔细建模也增加了人工成本。万亿模型的实际落地如何解决万亿模型训练的实际问题最大化利用硬件算力、同时降低训练成本。昇思MindSpore 针对上述三个挑战推出了一整套大模型训练解决方案作为一种通用的技术方案能够以较高效率实现万亿模型的训练。这三个挑战的解法分别是定义通信算子的反向算子、构建代价模型支持策略搜索、内置多种易用的并行技术。挑战一通信算子的自动微分为了实现通信算子的自动微分昇思MindSpore 定义了通信算子的反向算子。例如AllGather 的反向算子为 ReduceScatter。定义这些反向算子十分重要因为自动微分Auto-diff过程可以一次性地区分整个前向图而无须跳过任何算子——这也是为什么自动微分是自动并行Auto-parallel后面的必经一步。开发者无需手动在反向阶段插入 SEND/RECV 之类的梯度传递原语训练代码的编写难度随之大幅下降。挑战二代价模型与策略搜索针对第二个挑战昇思MindSpore 在同时考虑计算开销和通信开销的前提下建立一个代价模型来选择好的策略。为了快速地为复杂大图找到好策略提出了几种方法支持多图操作的算法将原始图转换成线性图降低搜索复杂度策略分离机制在保证返回解精度的同时有效缩小搜索空间。这些方法在工程实践上取得了可见的效果例如 ResNet50 在 8 台设备上搜索并行策略的时间在 1s 以内而返回的解决方案确实缩短了训练时间当模型较大类别数超过 128K时返回的解决方案与原始数据并行策略相比减少了大约 55% 的训练时间。需要说明的是这些数据源自本仓库课程文档的描述具体数值会随模型结构、集群配置与软件版本而变化。挑战三内置多级并行技术昇思MindSpore 内置了多种并行技术以易用的接口提供混合并行、流水线并行、异构训练和优化器并行等技术结合这些技术就能以较高的训练效率实现大模型训练。下面的代码展示了一个通过shard方法配置半自动并行的子模型示例class Submodel(nn.Cell): def __init__(self, shape): super().__init__() self.bn BatchNorm().shard(((4, 1), )) self.matmul MatMul().shard(((1, 1), (1, 4))) self.W Parameter(Tensor(shape), require_gradTrue) def construct(self, X): Y self.bn(X) Z self.matmul(Y, self.W) return Z这段代码中BatchNorm的输入张量按行切分为 4 份数据并行MatMul的权重W按列切分为 4 份模型并行二者之间的布局转换由框架自动插入张量重排布完成开发者无需感知具体通信细节。算子级别并行半自动并行与 shard 配置昇思MindSpore 支持开发者指定的高级策略配置称之为半自动并行semi-auto-parallel。如上文代码所示该子模型的结构为 BatchNorm 算子后跟一个 MatMul 算子广泛应用于 ResNet、ReID 等分类任务。在 BatchNorm 算子中输入张量 X 按行拆分为四部分数据可以并行效率非常高在 MatMul 算子中可学习参数权重 W 被分成四部分模型可以并行由于参数数量较多这部分的模型并行更有效。由于 BatchNorm 的输出布局与 MatMul 的输入布局不同框架会自动插入一个张量重排布该例中为 AllGather 与 ConCat这一过程对开发者完全透明开发者不必关注哪个设备运行了模型的哪个部分框架会自动安排。然而不同的模型结构在每个算子中具有不同大小的参数并且适用于不同的切分策略例如在以 Transformer 为代表的 NLP 子结构、以 ReID 为代表的图像分类子结构、以及 Wide Deep 类推荐模型子结构中第一算子配置为模型并行、后续算子配置为数据并行往往能获得更好的性能此时同样需要插入张量重排布该结构对比可参见本目录下的07MSParallel02.png示意图。在训练新模型时多次手动配置shard耗时耗力。这种情况下如果配置自动并行则不需要调用shard方法算法会自动找到一个有效的策略。例如当 ResNet 中的分类数量超过 130K 时算法返回的策略可以在 50ms 内完成一个迭代的训练而原始数据并行训练一次迭代超过 111ms数据同样来自本仓库课程文档。从框架架构看算子级并行属于编译优化MindCompiler层的能力昇思MindSpore 通过 函数变换接口中的shard指定函数输入/输出 Tensor 的分布策略配合基于 MindIR 的编译优化与自动微分机制将开发者声明式的高层策略翻译为带通信算子的分布式计算图可参考 昇思MindSpore关键特性 中对编程范式与函数变换接口的介绍。函数式算子切分动态图下的算子级并行动态图PyNative 模式语法更丰富、使用更灵活但目前昇思MindSpore 的动态图模式不支持自动并行的各种特性。借鉴 JAX 的pmap设计理念昇思MindSpore 设计了函数式算子切分功能支持在动态图模式下指定某一部分以图模式下的算子级并行方式执行。昇思MindSpore 的动态图模式下可以通过jit装饰符指定某一段以图模式编译执行。在前向执行的同时会将执行的算子、子图记录下来前向执行完毕后会对得到的整图进行自动微分得到反向图具体流程可参见 本目录下的 07MSParallel03.png 示意图。函数式算子切分沿用这一模式不同的是可以指定某一段在图模式的编译执行环节进行算子级模型并行。算子级并行通过将网络模型中每个算子涉及到的张量进行切分降低单个设备的内存消耗。这本质上利用了 MindSpore 基于源码转换的动静统一能力jit将 Python 函数编译为一张可调用的 MindSpore 图shard则对该函数内算子的输入/输出 Tensor 指定分布策略两者配合即可在动态图程序中按需点亮局部的高性能静态并行执行。流水线并行流水线Pipeline并行是将神经网络中的算子切分成多个阶段Stage再把阶段映射到不同的设备上使得不同设备去计算神经网络的不同部分。流水线并行适用于模型是线性图结构的场景能够降低模型训练过程中的通信量极大提升集群的训练性能。如图所示将 4 层 MatMul 的网络切分成 4 个阶段分布到 4 台设备上。正向计算时每台机器在算完本台机器上的 MatMul 之后将结果通过通信算子发送Send给下一台机器同时下一台机器通过通信算子接收Receive上一台机器的 MatMul 结果并开始计算本台机器上的 MatMul。反向计算时最后一台机器的梯度算完之后将结果发送给上一台机器同时上一台机器接收最后一台机器的梯度结果并开始计算本台机器的反向。微批次与流水线调度简单地将模型切分到多设备上并不会带来性能提升模型的线性结构会导致同一时刻只有一台设备在工作而其它设备在等待造成资源浪费。为了提升效率流水线并行进一步将小批次mini-batch切分成更细粒度的微批次micro-batch在微批次中采用流水线式的执行序。例如将小批次切分成 4 个微批次4 个微批次在 4 个组上执行形成流水线调度示意可参见 07MSParallel05.png。微批次的梯度汇聚后用来更新参数其中每台设备只存有并更新对应组的参数。流水线调度的思想与业界经典的 Gpipe / PipeDream 一脉相承Gpipe 采用先完整前向再完整反向的 F-then-B 策略并引入微批次以减少气泡PipeDream 则采用 1F1B 策略做完一个微批次的前向后立即做其反向进一步提升设备利用率详见 流水线并行章节。MindSpore 的内存优化调度昇思MindSpore 支持流水线并行并对执行序进行了调整以达到更优的内存管理在第 0 个 MicroBatch 的正向执行完后立即执行其反向这样做使得第 0 个 MicroBatch 的中间结果的内存得以更早地释放相较于先全部前向、再全部反向的调度进而确保内存使用的峰值更低调度对比可参见 07MSParallel06.png。MoE 并行混合专家Mixture of ExpertsMoE能够在引入较少计算量的同时扩充模型参数量。在模型中每一个 MoE 层都包含多个独立的 FFN前馈网络并包含一个路由装置Gate用于将每个输入数据分配到对应的 FFN。传统的数据并行和模型并行在处理 MoE 这种结构时存储和通信的效率都较为低下。一种高效的实现方式是将Gate 函数的计算并行化同时将多个 FFN 分配到每台设备上通常是均匀分配每台设备得到一个或多个 FFN每台设备负责的参数量因此有效下降。若训练数据经过 Gate 计算后路由到本台设备负责的 FFN那么数据直接传递给本台设备的 FFN如果路由到其他设备的 FFN那么会经过AllToAll 通信将训练数据发送到目的设备。同样地在经过 FFN 计算后数据需再次经过 AllToAll 通信把相应的结果路由回原设备当每条训练数据仅会路由到较少 FFN 时这种并行方式会非常有效它大大降低了内存使用同时产生的通信量也较小。这与仓库并行系列课程中 混合并行、ZeRO 数据并行 等章节讨论的通过切分降低单卡存储、用拓扑感知通信控制开销的思路相互印证共同构成大模型训练的多维并行手段。异构并行训练异构并行训练方法通过将占用内存较大的参数存储于 Host 内存解决单卡上无法存储完整模型的问题。其核心流程为通过分析图上算子的内存占用和计算密集度将内存消耗巨大或适合 CPU 逻辑处理的算子切分到 CPU 子图将内存消耗较小、计算密集型的算子切分到硬件加速器子图框架协同不同子图进行网络训练使得处于不同硬件且无依赖关系的子图能够并行执行架构示意可参见 07MSParallel08.png。这种 Host 内存与加速器协同的机制结合昇腾Ascend芯片整图下沉等能力见 昇思MindSpore关键特性 中的 Ascend 加速介绍目前昇思MindSpore 最多可以支持单机 8 卡训练和推理千亿模型极大降低训练所需的硬件资源。小结与思考昇思MindSpore 引入张量重排布Tensor Redistribution技术自动解决设备间张量布局不一致问题简化并行策略搜索并支持混合并行、流水线并行等高级并行技术。该框架通过自动微分技术定义了通信算子的反向算子如 AllGather 的反向算子为 ReduceScatter实现通信算子的自动微分从而简化并行训练过程。昇思MindSpore 支持算子级别并行和半自动并行配置允许开发者通过shard高级策略配置实现数据并行和模型并行同时自动插入必要的张量重排布操作提高并行效率。针对大模型训练昇思MindSpore 提供了如 MoE 并行、异构并行训练等解决方案优化存储和通信效率实现单机多卡训练千亿级参数模型的能力。对于希望进一步深入学习的读者可以在本仓库并行系列中继续阅读 数据并行、张量并行、流水线并行 与 混合并行 各章节结合 昇思MindSpore关键特性 中关于动静统一、函数变换与自动微分的介绍形成对大模型分布式训练全貌的系统认知。赞分享文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载相关推荐终极Tortoise-TTS分布式训练指南多GPU并行策略与优化技巧终极Tortoise TTS分布式训练指南多GPU并行策略与优化技巧 Tortoise TTS作为一款高质量多语音文本转语音系统其复杂的模型架构需要强大的计人工智能语音音频Apache MXNet模型并行与数据并行分布式训练策略对比Apache MXNet模型并行与数据并行分布式训练策略对比 引言分布式训练的双引擎 你是否曾遇到过训练大型深度学习模型时的困境单GPU内存不足、训练时间人工智能深度学习机器学习RAG 检索不再答非所问混合检索 智能重排序的四站流水线RAG 检索不再答非所问混合检索 智能重排序的四站流水线 给内部支持机器人提了个单“出差打车能报多少”机器人回了一段产品功能介绍。语料库里明明躺着完整示例工程上一篇基于 MCP 的 Loop Engineering 连接器实践最小权限、人类闸门与安全写入模式下一篇告别重复劳动用EZCard批量生成你的桌游卡牌创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考