ERASE：一种加速现代推荐系统训练的早期反向传播调度方法

# 用于加速现代推荐系统训练的早期反向传播调度方案
来源：https://arxiv.org/html/2608.18469
###### 摘要

轻量级代理模型无需反复训练前沿规模系统即可实现快速实验，但其小型计算核常导致现代加速器资源未充分利用。传统训练方法将前向与反向传播划分为独立阶段，这种低效性因此加剧——前一阶段的闲置容量无法被另一阶段的工作填充。我们将前向-前向（FF）算法的分离机制重新诠释为一种调度原语：基于局部目标，分离某模块的输出可移除下游梯度依赖，使得该模块的反向传播在其前向传播完成后即可就绪。*ERASE*系统在独立CUDA流上提前启动每个分离子图的反向传播，使其与后续前向工作重叠执行。轻量级Transformer上的执行轨迹展示了这种重叠及其限制：当计算核饱和设备时，将无法留出并发容量。在大规模点击率模型中，分离六个密集子架构使训练吞吐量提升达9.51%，同时保持归一化熵接近基准水平。

## 1引言

反向传播[21, 22]仍是跨视觉与语言[4, 26]、智能体学习与评估[18, 14, 28]、政治与社交网络推断[16, 24]以及自然科学[20, 27]领域训练现代神经网络的标准方法。其最具计算挑战性的应用包括大规模推荐与排序系统[2, 3]——这些系统需在超越单一训练任务数据流速度的持续数据流上反复重训练。因此训练吞吐量直接影响模型时效性与硬件资源负担[8]。

对于此类工作负载，吞吐量取决于加速器利用率与峰值算力性能[29]。集合通信、数据搬运与调度依赖关系可能导致容量闲置[17, 8]。传统反向模式微分按逆拓扑序遍历计算图，仅在整个前向传播完成后才启动反向传播。两阶段因此占据不相交的时间区间，导致任一阶段均无法利用另一阶段的闲置容量。

同样的利用率压力也出现在规模谱的另一端：前沿模型所需资源过多，无法支持快速重复实验，因此自动化研究依赖如NanoChat[11]这类轻量级代理——它在数量级更小的规模下保留了有用的实验信号。然而加速器规模并不会随模型缩小。小型GEMM操作暴露的并行度可能不足以接近峰值算力，即使没有启动间隙也会导致容量闲置；当间隙出现时，CPU调度计算开销可能进一步降低利用率。因此我们探索如何提升模型FLOPs利用率（MFU），特别是当模型设计过小无法饱和设备的情况。

与此同时，前向-前向（FF）算法对每个模块基于局部目标进行训练，并向其后继传递*分离的*激活值[9]。这将端到端计算图划分为模块级子图：梯度不跨越模块边界，激活值无需为全局反向传播保留。对调度而言尤其重要的是：模块的反向传播在前向传播与局部目标完成后即可就绪，无需等待终端损失或后续模块。因此它可以与后续前向传播重叠，实现传统前向-后向调度所放弃的工作。

我们将这一观察转化为*ERASE*系统（第2节），它在每个分离子图的前向传播返回后立即调度其反向传播。实验（第3节）首先在NanoChat上验证了预期的重叠及其饱和限制，随后在排序模型上测量吞吐量提升。第4节分析了异步调度引入的错位现象。

### 1.1相关工作

##### 无反向传播学习

无反向传播方法用模块局部目标替代全局反向传播。NoProp和DiffusionBlocks为每个模块附加辅助去噪变量[13, 25]，借鉴扩散模型[10, 23]的训练信号；这允许分离激活值与模块局部更新。早期关于前向-前向与逐层学习的工作同样表明，局部目标可在不跨越模块边界的情况下训练出具有竞争力的模型[9, 15]。

##### CUDA图与执行级加速

CUDA图捕获并回放计算核序列，减少小计算核工作负载中的每核CPU启动开销[19, 6]；近期的编译器工作使捕获在PyTorch中更稳健[7]。ERASE将这种执行机制与FF式分离相结合，在保持常规损失与各模块内精确梯度的同时创建独立计算单元。早期反向传播在每个分离单元的前向传播返回后立即在独立CUDA流上启动，而CUDA图则确保跨排名的异步启动顺序一致。

## 2无反向传播算法

### 2.1前向-前向算法

##### 数据依赖阻断

对于输入嵌入$\bm{x}\in\mathbb{R}^{p}$，考虑包含$B$个模块的网络，其中$\bm{z}_0=\bm{x}$，激活值$\bm{z}_b$和参数$\bm{\theta}_b$。FF算法[9]计算：
$$\bm{z}_b=f_b\big(\operatorname{sg}[\bm{z}_{b-1}];\bm{\theta}_b\big),\qquad b=1,\dots,B,(1)$$
其中$\operatorname{sg}[\cdot]$是梯度停止算子：前向传播时为恒等映射，反向传播时雅可比矩阵为零。对未分离的$\bm{z}_b$施加局部损失仅更新$\bm{\theta}_b$，例如可对“良好度”分数$G_b=\sum_j z_{bj}^2$应用二元交叉熵。因此$\nabla_{\bm{\theta}_b}$仅依赖于相邻切割点间的子图，在模块前向传播与局部损失完成后即就绪，无需等待终端损失评估。ERASE利用了分离带来的这一调度特性；图1展示了由此产生的计算图。

图1：FF式分离示意图。每个模块接收前一模块激活值的分离副本，并贡献自身的局部损失。梯度不跨越模块边界。

### 2.2*ERASE*中的早期反向传播：调度、流与确定性

分离暴露了独立性但不改变调度：对聚合损失
$$\mathcal{L}=\mathcal{L}_{\text{main}}+\sum_{i}\lambda_{i}\,\ell_{i}(2)$$
的单次反向传播调用仍在训练步末尾执行。ERASE改为在每个分离子图的前向传播返回后立即启动其反向传播，在保留未分离部分执行常规步末反向传播的同时累积其参数梯度。由于反向传播操作否则会与默认CUDA流上的后续前向工作串行化，ERASE使用独立流与事件强制实施剩余依赖关系。这使得模块$b$的反向传播能与模块$b+1$的前向传播重叠。

CPU调度引入了另一权衡。主线程早期反向传播会阻塞进一步调度，直至自动求导调用入队。线程池执行器可避免此停滞，但可能使跨排名的计算核与集合顺序产生差异，形成我们称为“错位”的落后节点。将受影响的子图捕获为CUDA图可固定启动顺序。第4节比较了阻塞与非阻塞变体。

## 3实验

作为小规模可行性验证，在单张A100 GPU上，对使用MNIST数据集[5]的3层MLP应用FF式分离，使计算匹配的反向传播时间减少58%（23.5至9.9毫秒），批次总时间减少30%（41.2至28.8毫秒）。随后我们在NanoChat[11]上评估*ERASE*，使用单张A100 GPU验证重叠效果；并在大规模推荐模型的轻量版本上使用八张H100 GPU测量吞吐量。

### 3.1 NanoChat

为在推荐模型实验前验证预期重叠，我们使用两个分离点将NanoChat的多头注意力栈划分为三个子图。图2展示了单批次执行轨迹。

参考图注图2：采用两个分离点与三个流的NanoChat单批次轨迹。一个子图的反向工作与下一个子图的前向工作并行运行。该轨迹展示了执行模式而非速度提升：一个子图的反向工作与下一个子图的前向工作重叠，如第2.2节预测。例外是融合多头注意力反向传播核（图2中的*fm*），它单独运行。这是资源限制而非依赖问题：通过融合注意力矩阵乘法与softmax并避免序列间注意力矩阵的物化，该计算核占满所有流式多处理器，未留下并发工作容量。

此限制强化了第1节的前提：重叠仅在硬件尚未饱和时有效。减小隐藏宽度、深度或批次尺寸可缩小计算核而不缩小加速器，留下小型计算核可共享但设备填满式计算核无法使用的资源。轻量级代理主要包含前者，使其成为早期反向传播的天然目标。第3.2节量化了其在推荐模型上的收益。

### 3.2点击率模型

#### 3.2.1设置

后续实验采用包含六个分离子架构的点击率（CTR）模型，其中四个被分配独立CUDA流与主流并行。吞吐量以每秒处理的训练样本数（QPS）衡量，通过热身后的步样本p90（越高越好）汇总；与延迟p90不同，这属于分布中的快速端。质量指标为归一化熵（NE），即模型交叉熵除以常量预测器的交叉熵；越低越好。

#### 3.2.2 CTR模型上的早期反向传播

ERASE在每个分离子架构的前向传播完成后立即通过线程池执行器启动其反向传播。由于工作线程可能引入跨排名启动顺序不确定性，我们将部分子架构捕获为CUDA图以固定其顺序并降低CPU调度开销。此非阻塞配置（表1第1行）提升吞吐量7.38%，NE差距约1.38%（图3）。第4节探讨了其FUP=False设置并与阻塞调度进行比较。

参考图注图3：非阻塞早期反向传播变体（CUDA图）与基准的NE差距曲线，约1.38%。

## 4注意事项

表1报告了调度与FUP（find_unused_parameters）消融实验。*阻塞*模式从主线程启动早期反向传播，并在自动求导入队期间阻塞CPU调度；FUP控制参与梯度同步的参数。

表1：CTR模型上的吞吐量，p90 QPS阻塞模式具有确定性但阻塞CPU：FUP=False仅获得0.37%提升（第3行），而FUP=True通过减少最终聚合反向传播的参数获得5.26%提升（第2行）。CUDA图通过非阻塞调度与FUP=False实现7.38%增益（第1行）。综合结果表明，确定性、排名同步的集合顺序具有实用价值，且FUP或CUDA图均可提供此保证。

## 5结论

*ERASE*将FF式分离重新用作调度原语：切断模块间依赖使每个反向传播在前向传播后即可就绪，在保持子图内精确梯度的同时实现跨流重叠，无需FF的“良好度”目标。在CTR模型上，ERASE在确定性跨排名集合顺序下将p90 QPS提升5-9%，伴随轻微NE差距。

## 参考文献

（保留原始参考文献格式与链接）