arXiv:2512.08160cs.LGcs.AI2025-12被引 1

提出可精确控制梯度延迟的分阶段训练框架,降低显存占用并提升训练效率。

LayerPipe2: Multistage Pipelining and Weight Recompute via Improved Exponential Moving Average for Training Neural Networks

  • 基于指数移动平均重构历史权重,避免显式存储
  • 内层延迟少、外层延迟长,延迟量由下游层数决定
  • 为流水线训练提供理论依据,适合大规模模型训练

在先前工作LayerPipe中,我们提出通过重叠前向与反向计算来加速卷积网络、全连接网络及脉冲神经网络的训练。然而,尽管实验成功,各层应引入多少梯度延迟以达到理想流水线效果仍缺乏理论解释。本文通过变量延迟梯度自适应与重定时分析,形式化推导了LayerPipe。我们发现延迟可合法插入的位置,并证明延迟量直接由网络结构决定:内层所需延迟较少,外层需更长延迟。当每层均流水时,延迟量仅取决于剩余下游阶段数;当分组流水时,组内各层共享相同延迟分配。这些洞察不仅解释了以往观测到的调度模式,还揭示了一个常被忽视的问题——流水线隐含需要存储历史权重。为此,我们提出一种面向流水线的移动平均机制,通过重建而非显式存储过去状态,有效降低内存开销,同时保持流水训练的精度保障。该框架提供了构建LayerPipe架构的理论指导,可预测延迟需求并缓解存储瓶颈,实现可控通信-计算权衡的大规模流水线训练。

原文摘要 · Abstract (English)

In our prior work, LayerPipe, we had introduced an approach to accelerate training of convolutional, fully connected, and spiking neural networks by overlapping forward and backward computation. However, despite empirical success, a principled understanding of how much gradient delay needs to be introduced at each layer to achieve desired level of pipelining was not addressed. This paper, LayerPipe2, fills that gap by formally deriving LayerPipe using variable delayed gradient adaptation and retiming. We identify where delays may be legally inserted and show that the required amount of delay follows directly from the network structure where inner layers require fewer delays and outer layers require longer delays. When pipelining is applied at every layer, the amount of delay depends only on the number of remaining downstream stages. When layers are pipelined in groups, all layers in the group share the same assignment of delays. These insights not only explain previously observed scheduling patterns but also expose an often overlooked challenge that pipelining implicitly requires storage of historical weights. We overcome this storage bottleneck by developing a pipeline--aware moving average that reconstructs the required past states rather than storing them explicitly. This reduces memory cost without sacrificing the accuracy guarantees that makes pipelined learning viable. The result is a principled framework that illustrates how to construct LayerPipe architectures, predicts their delay requirements, and mitigates their storage burden, thereby enabling scalable pipelined training with controlled communication computation tradeoffs.

流水线训练梯度延迟显存优化深度学习

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。