在大模型训练与推理的实践中,并行优化是提升效率、降低资源消耗的核心手段。MindSpeed作为面向大模型的高效训练框架,通过多种并行策略的协同设计,实现了模型规模扩展与计算资源利用率的平衡,其核心优化思路可概括为“切分权重与激活值、分层并行、专家并行”三大方向。
一、PP(流水线并行)与VPP(虚拟流水线并行):切分模型,释放内存与性能潜力
- PP(Pipeline Parallelism):核心是切分模型层权重,将模型的不同层分配到不同设备(如GPU)上,使每个设备仅需存储和处理部分模型参数,从而降低单设备内存压力。例如,将大型Transformer的多层结构按层拆分,Device 1处理前几层,Device 2处理后几层,形成流水线式的计算流程(Forward Pass + Backward Pass交替进行)。
- VPP(Virtual Pipeline Parallelism):在PP基础上进一步细化切分粒度,将单个PP Stage再拆分为多个子阶段,减少设备间的空泡率(即设备等待上游/下游数据的时间)。通过“Assign multiple stages to each device”(为每台设备分配多个子阶段),让设备在计算时不闲置,最大化硬件利用率,尤其在大模型层数极多的场景下,性能提升显著。
二、分布式优化器(ZeRO - 1)与Param/grad Buffer Sharding:分散状态,降低内存与计算时间
- 分布式优化器(ZeRO - 1):开启后,将优化器的状态(如动量、方差等)均匀分散到各数据并行(DP)的设备(DP)中。传统优化器状态集中存储会导致单设备内存爆炸,ZeRO - 1通过“状态拆分”,让每个DP设备仅存储部分优化器状态,大幅降低内存占用;同时,计算时各设备独立更新自身负责的状态,减少跨设备通信开销,缩短整体计算时间。
- Param and grad Buffer Sharding:将模型参数(Param)和梯度(grad)的缓冲区按“world(全局)、local(局部)、param(参数)”维度分片。结合图例中的“dp rank”(数据并行秩),不同DP rank的设备仅处理和存储部分参数/梯度,进一步降低单设备内存压力,同时通过“分而治之”的计算方式,提升并行计算效率。
三、TP(张量并行)+ SP(序列并行)与专家并行:切分权重/激活,适配复杂模型结构
- TP + SP(Tensor Parallelism + Sequence Parallelism):TP侧重切分模型层的权重(如将矩阵乘法按维度拆分到不同设备),SP则切分序列维度的激活值(如将长序列的输入按长度拆分)。二者结合,既降低了单设备的权重存储压力,又通过序列切分减少了单设备处理的序列长度,适合长文本、多模态等序列密集型任务,实现“权重与激活值双重切分,内存与计算双优化”。
- 专家并行(Expert Parallelism):针对混合专家模型(MoE),将专家网络切分到不同设备上。每个专家(Expert 1/2/3/4…)独立部署在不同设备,通过门控(Gating)机制动态分配输入到对应专家,降低单设备的内存与计算负载(仅需处理部分专家的计算)。同时,结合“Un - Permutation”(反置换)操作,解决专家输出后的数据重组问题,保障并行计算的正确性。
四、并行优化的协同价值
MindSpeed的并行优化并非单一策略的使用,而是多种并行方式的协同组合:PP/VPP负责模型层的纵向切分,ZeRO - 1与Buffer Sharding负责优化器状态的横向分散,TP/SP负责权重与激活的多维切分,专家并行负责MoE模型的专家级拆分。这种“分层 + 多维 + 模型适配”的并行设计,让大模型在训练时能充分利用硬件资源(多GPU、多节点),在降低内存消耗的同时,大幅提升计算效率,为大模型从“实验室可行”走向“工业级落地”提供了关键技术支撑。