arXiv:2410.07531cs.ARcs.AI2024-10

通过让随机数生成与GEMM重叠,显著降低大模型注意力中的丢弃开销。

Reducing the Cost of Dropout in Flash-Attention by Hiding RNG with GEMM

  • 将随机数生成与前序GEMM层并行执行,避开硬件瓶颈
  • 在Llama3+FP8+GH100上实现1.26倍训练加速
  • 适用于多种模型和硬件,可直接集成到现有系统

Dropout 在启用时会显著影响 Flash-Attention 的性能,从而延长大语言模型(LLMs)的端到端训练时间。主要原因是随机数生成(RNG)阶段的延迟。现有优化将 RNG 融入 Flash-Attention 内核,但由于 RNG 与 Attention 共享底层硬件瓶颈,融合难以有效隐藏延迟。本文提出将 RNG 与前序 GEMM 层重叠执行,利用两者资源需求和硬件瓶颈不同,实现无干扰并行。我们构建了细粒度性能分析模型,经硅片实测验证,在单个 Transformer Block(包含多头注意力和前馈层)上,相比串行实现提升 1.26 倍速度,比当前最优融合方案快 1.22 倍,适用于 Llama3 在 GH100 GPU 上使用 FP8 精度的场景。由于内核模式规律性强,该方法可推广至不同模型架构、软件实现和硬件配置。

原文摘要 · Abstract (English)

Dropout, a network operator, when enabled is likely to dramatically impact the performance of Flash-Attention, which in turn increases the end-to-end training time of Large-Language-Models (LLMs). The main contributor to such performance degradation is the Random Number Generation (RNG) phase. The state-of-the-art optimization is to fuse RNG into the Flash-Attention kernel. However, while RNG and Attention do not compete on compute or memory resources, they are bounded on the same lower-level architecture bottlenecks. Fusion can hardly hide RNG latency within the Attention kernel. We propose overlapping RNG with previous GEMM layers in the network to hide RNG latency and improve end-to-end performance. RNG and GEMM have distinct resource requirements and hardware bottlenecks, so they can run together without compromising each other's performance. We propose a fine-grained analytical performance model that analyzes low-level architecture resource utilization to evaluate RNG-GEMM overlapping performance benefits. This model, cross-validated by silicon results, shows 1.26x speedup for overlapping RNG and GEMM layers over a sequential implementation on one Transformer Block (one LLM layer including multi-head attention and feed-forward layers), and 1.22x over state-of-the-art fusion implementation, for Llama3 on GH100 GPUs with FP8 precision. Because the kernel patterns are regular, the findings of the shared bottlenecks, as well as the achievable performance benefits, can be generalized to different model architectures, software implementations and hardware configurations.

Flash-AttentionDropout优化GEMM并行大模型训练

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