迈入大模型时代的深度学习：使用 Flash Attention 技术让 Transformer 起飞

Transformer 是 ChatGPT 等大语言模型的核心技术之一，而注意力机制是其的关键部分。但是，标准的注意力实现具有二次时间和内存复杂度，使其在长序列任务上变慢并消耗大量显存。这限制了 Transformer 能够建模的上下文长度，同时使得大尺度模型的训练和推理时间延长。

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness [1] 提出了一种称为 Flash Attention 的新注意力算法，该算法减少了注意力机制所需的内存访问次数。这使得注意力运算更快速且内存高效，同时仍保持数值精确性。

                        图｜FlashAttention 的运行机制 [1]

Flash Attention 运用了两种主要技术:
1. 分块 - 它将输入划分成块，并单独处理每个块的注意力。

2. 重计算 - 它只存储足够的信息，以便在反向传播期间重新计算注意力矩阵，而不是存储整个矩阵。这减少了内存使用。

与标准注意力相比，Flash Attention 通过减少内存访问次数从而提高训练速度并降低显存占用量。相比于基准，Flash Attention在GPT-2和BERT等 Transformer 模型上的训练速度提升可达 3 倍以上。

                        图｜FlashAttention 模型训练速度对比 [1]

---

FlashAttention Notebook

Flash Attention 提供了一种新型高效的 Transformer 加速技术，通过优化 IO 访问提高了计算速度并降低显存占有。同时通过 Flash Attention 训练获得的模型能够无感知迁移到标准 Transformer 框架中而不牺牲运算精度。

本期 NoteBook 中详细介绍了 Flash Attention 技术的基本原理，同时提供了 GPT 模型的上手案例以便于读者更加容易理解与使用 Flash Attention 技术。

点击下方图片，访问本期 Notebook：

Reference:
[1] https://doi.org/10.48550/arXiv.2205.14135

---

欢迎来 Notebook 案例广场，获取更多有意思的notebook实践～
感兴趣的同学也可以点击原文查看。