自定义空操作层

背景与挑战

在神经网络训练过程中，初始层的嵌入（Embedding）操作及输出层的对数几率（Logits）计算通常属于计算密集型任务，这可能对整个网络的运行效率产生显著影响。

• 嵌入层（Embedding Layer）：在处理文本或分类数据时，嵌入层将高维稀疏特征转换为低维稠密向量表示。此过程涉及索引查找和潜在的大规模矩阵乘法，特别是在自然语言处理应用中，词汇表可能包含数十万乃至数百万词条。高维度的查找与转换操作会消耗大量计算资源。
• 对数几率层（Logits Layer）：位于网络末端的对数几率层通常是一个全连接层，其功能是将最后一层的隐藏状态映射到输出空间，为后续损失函数计算提供未归一化的预测值。如果分类任务具有大量类别，那么该层的权重矩阵将非常庞大，导致矩阵乘法运算成为性能瓶颈。

上述操作的计算复杂度随着输入特征数量和类别数量的增加而上升，可能导致训练速度降低，并且在计算资源有限的环境中形成性能瓶颈。

解决方案

为应对上述挑战，我们引入了“自定义空操作层”功能，允许用户通过指定特定层为“空操作层”（No-Op Layers）来动态调整模型在训练流水线中的计算负载。此机制有助于在多个计算节点间更均匀地分配工作负载，从而优化整体计算资源的利用。

使用场景

在开启流水线并行时，当用户遇到由于计算负载不均导致的性能瓶颈时，此功能尤为适用。通过对计算密集型任务进行重新分配，可以有效减少流水线中的空闲时间（即空泡），从而提高系统的吞吐量和效率。

使用方法

要启用此功能，用户需通过命令行参数设置目标层为无操作层。例如，原模型共126层，该模型配置如下参数，即执行实际计算的层有126层。

--num-layers 126

若在该模型首和尾各自添加1层空层，则该模型配置如下参数，表示总共128层。首尾层（即第0层和第127层，层数从0开始计数）为不执行实际计算的空操作层，中间的126层为执行实际计算的层。

--num-layers 128 
--noop-layers 0,127 

使用该特性增加空层后总层数发生变化，需要根据包含空层的总层数重新调整流水线（虚拟流水线）的配置。

使用效果

通过实施自定义增加无操作层的策略，预期能够显著减少流水线中的空泡现象，从而优化计算流程并提升系统性能。这不仅有助于加速模型训练过程，还能最大化硬件资源的利用率。

• 减少空泡现象：通过将计算密集型任务重新分配到流水线的不同阶段，减少了因某些阶段计算负载过高而导致的空闲时间，提高了整个系统的吞吐量。
• 优化计算流程：确保计算和通信资源被更加充分地利用，避免了因同步障碍造成的闲置浪费，使得每个计算单元都能高效工作。

根据实际测试结果，在类DeepSeekV2百亿参数级别的MoE模型中，由于终端层的对数几率层等导致的计算负载不均所引起的性能瓶颈，使用自定义空操作层可以获得超过10%的端到端训练性能提升。