自定义空操作层

背景与挑战

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

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

解决方案

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

使用场景

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

使用方法

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

--num-layers 126

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

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

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

使用效果

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

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