FlashAttention-2: Attention nhanh hơn nhờ song song hóa và phân chia công việc tốt hơn

• Các mô hình ngôn ngữ có ngữ cảnh dài hơn như GPT-4(32k), MPT(65k), Calude(100k) đang xuất hiện
• Việc mở rộng độ dài ngữ cảnh của Transformer là khó vì yêu cầu runtime & memory tăng theo lũy thừa bậc bốn
• FlashAttention ra mắt năm ngoái đã được sử dụng ở nhiều nơi nhờ giảm mức dùng bộ nhớ và tăng tốc độ attention
• Ngay khi ra mắt, nó đã nhanh hơn 2-4 lần, nhưng vẫn còn dư địa để cải thiện. So với phép nhân ma trận tối ưu hóa (GEMM), nó vẫn chưa đủ nhanh và về mặt lý thuyết chỉ đạt 25-40% FLOPs/s tối đa (tối đa 124 TFLOPs/s trên GPU A100)
• FlashAttention-2 nhanh gấp 2 lần phiên bản trước và cung cấp hiệu năng tối đa 230 TFLOP/s trên GPU A100
• Trong huấn luyện mô hình ngôn ngữ kiểu GPT, nó đã đạt tới 225 TFLOPS (72% mức tận dụng FLOP của mô hình)
• Đã điều chỉnh thuật toán để giảm non-matmul FLOPs
• Song song hóa tốt hơn, thay đổi cách phân chia công việc trong mỗi thread block
• Mở rộng số lượng Head Dimensions từ 128 lên 256