Reproduce “DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training” on OpenResearch

About this paper

Xinwei Qiang, Hongmin Chen, Shixuan Sun, Jingwen Leng, Xin Liu, Minyi Guo

Researchers from Shanghai Jiao Tong University and ByteDance Seed introduced DASH, a framework designed to improve the efficiency of deterministic attention backward passes in large language models. DASH, which optimizes computation schedules and accumulation orders, achieved up to 1.28x speedup for deterministic attention backward passes compared to FlashAttention-3's deterministic mode, ensuring bitwise identical gradient outcomes for reproducible LLM training.

Key insights

• The primary cause of performance degradation in deterministic attention is suboptimal scheduling that creates a misalignment between tile execution and enforced accumulation ordering, rather than the serialization itself.
• Optimizing deterministic attention requires a co-design approach that simultaneously considers the computation schedule and the accumulation order within the GPU.
• Theoretical optimality in scheduling can be constrained by practical hardware limitations, such as register pressure and L2 cache latency, leading simpler heuristics to sometimes outperform more complex, theoretically optimal algorithms on current architectures.

Summary by alphaXiv.