提升大模型训练确定性下的吞吐率，解决梯度累积导致的性能损失问题。
DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training
- 将确定性注意力反向传播建模为有向无环图调度问题，优化计算与梯度合并顺序。
- 在H800 GPU上实现最高1.28倍的反向传播吞吐提升，缩小确定性与非确定性性能差距。
- 提出两种互补调度策略，适合追求可复现性又需高效训练的LLM研究者使用。
确定性对大语言模型（LLM）训练的可复现性至关重要，但通常带来显著性能开销。在广泛使用的FlashAttention-3等注意力实现中，确定性反向传播相比非确定性版本吞吐率下降高达37.9%，主要因梯度累积操作必须串行化以保证数值一致性。该性能损耗源于计算与梯度归约阶段调度不当，造成硬件严重闲置。为此，我们将确定性注意力的反向传播建模为有向无环图（DAG）上的调度问题，并推导出最小化关键路径长度的调度方案。基于此，我们提出DASH（Deterministic Attention Scheduling for High-Throughput），包含两种互补调度策略：(i) 降序查询块遍历（Descending Q-Tile Iteration），通过逆序查询块访问减少因果注意力中的流水线停顿；(ii) 移位调度（Shift Scheduling），在我们的DAG模型下理论最优，可降低全注意力与因果掩码下的流水线停顿。在NVIDIA H800 GPU上的实证评估表明，DASH显著缩小了确定性注意力的性能差距，所提策略使注意力反向传播吞吐率最高提升1.28倍，大幅推动可复现性训练的效率。代码已开源于 https://github.com/SJTU-Liquid/deterministic-FA3。
原文摘要 · Abstract (English)
Determinism is indispensable for reproducibility in large language model (LLM) training, yet it often exacts a steep performance cost. In widely used attention implementations such as FlashAttention-3, the deterministic backward pass can incur up to a 37.9% throughput reduction relative to its non-deterministic counterpart, primarily because gradient accumulation operations must be serialized to guarantee numerical consistency. This performance loss stems from suboptimal scheduling of compute and gradient-reduction phases, leading to significant hardware underutilization. To address this challenge, we formulate the backward pass of deterministic attention as a scheduling problem on a Directed Acyclic Graph (DAG) and derive schedules that minimize the critical path length. Building on this formulation, we present DASH (Deterministic Attention Scheduling for High-Throughput), which encapsulates two complementary scheduling strategies: (i) Descending Q-Tile Iteration, a reversed query-block traversal that shrinks pipeline stalls in causal attention, and (ii) Shift Scheduling, a theoretically optimal schedule within our DAG model that reduces pipeline stalls for both full and causal masks. Our empirical evaluations on NVIDIA H800 GPUs demonstrate that DASH narrows the performance gap of deterministic attention. The proposed strategies improve the throughput of the attention backward pass by up to 1.28$\times$ compared to the baseline, significantly advancing the efficiency of reproducible LLM training. Our code is open-sourced at https://github.com/SJTU-Liquid/deterministic-FA3.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。