Source-linked AI summary
ERASE: EaRly bAckpropagation SchEdule for Faster Training of Modern Recommendation Systems
Ergan Shang, Flavio Sales Truzzi
TL;DR
Lightweight recommendation proxies can underuse accelerators because conventional training keeps forward and backward passes in separate phases. ERASE uses FF-style detachment to launch backward work early, improving p90 QPS by 5–9% on a CTR model with a small normalized-entropy gap under deterministic collective ordering.
Problem
Lightweight proxy models and disjoint forward/backward scheduling can leave accelerator capacity unused, limiting training throughput for rapidly retrained recommendation systems.
Method
ERASE repurposes FF-style detachment to make each subgraph’s backward pass ready after its forward pass, enabling separate-stream overlap with subsequent forward work.
Results
5–9% p90 QPS improvement on a CTR model is achieved with a small normalized-entropy gap under deterministic cross-rank collective order.
Takeaways & Limitations
ERASE shows that early backward scheduling can improve throughput for lightweight recommendation models when kernels leave hardware capacity available for overlap.
Takeaways & Limitations
Non-blocking dispatch can introduce cross-rank launch-order nondeterminism, while blocking dispatch can stall the CPU; CUDA Graph capture fixes launch order.
Abstract
from arXiv · showhide
Lightweight proxy models enable rapid experimentation without repeatedly training frontier-scale systems, but their small kernels often leave modern accelerators underutilized. Conventional training compounds this inefficiency by scheduling the forward and backward passes as disjoint phases, so spare capacity in one cannot be filled by work from the other. We reinterpret the detachment mechanism of Forward-Forward (FF) as a scheduling primitive: given a local objective, detaching a block's output removes downstream gradient dependencies, making its backward pass ready when its forward pass finishes. ERASE launches each detached subgraph's backward pass early on a separate CUDA stream, overlapping it with subsequent forward work. Execution trace on a lightweight transformer demonstrates this overlap and its limit: a kernel that saturates the device leaves no capacity for concurrency. On a large-scale click-through-rate model, detaching six dense subarchitectures improves training throughput by up to $9.51\%$ while keeping normalized entropy close to the baseline.
1 Introduction
ERASE addresses accelerator underutilization in lightweight recommendation-model proxies by turning FF-style detachment into an early-backward scheduling primitive. Detached subgraphs can launch backward work on separate CUDA streams while later forward passes proceed, subject to device saturation limits.
- Motivation: Accelerator utilization, not peak arithmetic performance alone, limits throughput when communication, data movement, and scheduling dependencies leave capacity unused.Conventional reverse-mode differentiation keeps forward and backward execution in disjoint intervals.
- Motivation: Lightweight proxy models enable rapid experimentation but expose insufficient parallelism through small GEMMs, leaving accelerator capacity unused.The accelerator does not shrink with the model, so smaller kernels may fail to approach peak arithmetic throughput.
- Scheduling primitive: FF-style detached activations partition the graph into block-local subgraphs whose backward passes become ready after their local forward passes and objectives complete.Gradients do not cross block boundaries, eliminating the need to wait for the terminal loss or later blocks.
- ERASE: ERASE dispatches each detached subgraph’s backward pass immediately after its forward pass returns, overlapping backward work with subsequent forward passes on separate CUDA streams.This combines FF-style detachment with CUDA execution mechanisms that create independent pieces while preserving usual loss and exact within-piece gradients.
- Evaluation: The intended overlap is demonstrated on NanoChat and its throughput benefits are measured on a ranking model, with asynchronous dispatch introducing execution staggering.The experiments also examine the saturation limit of overlap.
2 Backprop Free Algorithm
ERASE turns FF-style detachment into an early-backward scheduling mechanism: detached block gradients become locally ready after each forward pass, then execute on separate CUDA streams alongside later forward work. CUDA Graph capture addresses CPU-dispatch stalls and cross-rank staggering caused by non-blocking launches.
- Data dependency shutdown: Each detached block’s parameter gradient depends only on its local subgraph and becomes ready after its forward pass and local loss.This readiness occurs before evaluating the terminal loss.
- Data dependency shutdown: FF detaches each block’s input activation and assigns a local loss, so no gradient crosses block boundaries.The stop-gradient operator is identity in the forward pass but has zero Jacobian during backpropagation.
- Early-backward scheduling: ERASE launches each detached subgraph’s backward pass immediately after its forward pass, while the undetached remainder uses the usual end-of-step backward.Detachment alone does not alter scheduling; ERASE explicitly starts early backward work.
- Early-backward scheduling: Separate CUDA streams and events let early backward operations overlap subsequent forward work while preserving remaining dependencies.Backward operations would otherwise serialize with later forward work on the default CUDA stream.
- Dispatch tradeoffs: CUDA Graph capture fixes launch order when ThreadPoolExecutor dispatch avoids blocking but causes kernel and collective staggering across ranks.Main-thread dispatch blocks further work until autograd is enqueued, whereas thread-pool dispatch can create stragglers.
3 Experiments
Experiments show that early backward reduces compute-matched training time, creates the intended forward–backward overlap, and improves CTR-model throughput, while device-saturating kernels limit concurrency. On the CTR model, the non-blocking CUDA-Graph configuration improves throughput by 7.38% with an approximately 1.38% normalized-entropy gap.
- Small-scale sanity check: 58% lower compute-matched backward time and 30% lower total batch time were measured on a 30-layer MNIST MLP using one A100 GPU.Backward time fell from 23.5 to 9.9 ms, while total batch time fell from 41.2 to 28.8 ms.
- NanoChat overlap demonstration: NanoChat’s execution trace shows one subgraph’s backward work overlapping the next subgraph’s forward work across three streams.The trace uses two detachment points partitioning the multi-head-attention stack into three subgraphs.
- NanoChat overlap demonstration: The fused multi-head attention backward kernel runs alone because it occupies every streaming multiprocessor, making the limitation a resource constraint rather than a dependency.Fusion combines attention matrix multiplications with softmax while avoiding materialization of the sequence-by-sequence attention matrix.
- Implication for lightweight proxies: Lightweight proxies are natural targets for early backward because smaller kernels leave accelerator capacity available for sharing, unlike device-filling kernels.Reducing hidden width, depth, or batch size shrinks kernels without shrinking the accelerator.
- CTR-model evaluation: 7.38% higher throughput was achieved with an NE gap of approximately 1.38% using six detached CTR subarchitectures, CUDA streams, and CUDA Graphs.Backward launches begin when each forward completes; CUDA Graphs reduce CPU dispatch overhead and stabilize cross-rank launch order.
4 Caveats
The caveats show that blocking dispatch can stall CPU work, while FUP and CUDA Graphs improve throughput under different dispatch and synchronization configurations. Deterministic, rank-synchronized collective order appears useful, with either FUP or CUDA Graphs providing it.
- Dispatch and synchronization: Blocking launches early backward from the main thread and stalls CPU dispatch during autograd enqueue.FUP controls which parameters participate in gradient synchronization.
- Dispatch and synchronization: 0.37% gain occurs with FUP=False, compared with 5.26% with FUP=True by reducing parameters in the final aggregate backward.Both results use blocking dispatch; FUP=True reduces the parameters participating in the final aggregate backward.
- Dispatch and synchronization: 7.38% gain comes from CUDA Graphs with non-blocking dispatch and FUP=False.Together with FUP, CUDA Graphs can provide deterministic, rank-synchronized collective order.
5 Conclusion
ERASE repurposes FF-style detachment as a scheduling primitive, enabling separate-stream overlap while preserving exact within-subgraph gradients without FF’s goodness objective. On a CTR model, it improves p90 QPS by 5–9% with a small NE gap under deterministic cross-rank collective order.
- 5 Conclusion: ERASE cuts inter-block dependencies so each backward pass becomes ready after its forward pass, enabling overlap on separate streams.This repurposes FF-style detachment as a scheduling primitive.
- 5 Conclusion: ERASE preserves exact within-subgraph gradients without using FF’s goodness objective.
- 5 Conclusion: 5–9% p90 QPS improvement is achieved on a CTR model, with a small NE gap under deterministic cross-rank collective order.