Transformer 是 ChatGPT 等大语言模型的核心技术之一，而 注意力机制 是其的要害局部。然而，规范的注意力实现具备二次工夫和内存复杂度，使其在长序列工作上变慢并耗费大量显存。这限度了 Transformer 可能建模的上下文长度，同时使得大尺度模型的训练和推理工夫缩短。
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness [1] 提出了一种称为 Flash Attention 的新注意力算法，该算法缩小了注意力机制所需的内存拜访次数。这使得注意力运算更疾速且内存高效，同时仍放弃数值精确性。
图|FlashAttention 的运行机制 [1]
Flash Attention 使用了两种次要技术:
1. 分块 – 它将输出划分成块，并独自解决每个块的注意力。
- 重计算 – 它只存储足够的信息，以便在反向流传期间从新计算注意力矩阵，而不是存储整个矩阵。这缩小了内存应用。
与规范注意力相比，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 实际～
感兴趣的同学也能够点击原文查看。