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實踐～
感興趣的同學也可以點擊原文查看。