【大模型上下文长度扩展】FlashAttention-2：比1代加速1.29倍、GPU利用率从55%上升到72%

2024-02-07 02:12
文章标签 模型 gpu 加速 扩展 长度 上下文 利用率 上升 flashattention 55% 72% 1.29

本文主要是介绍【大模型上下文长度扩展】FlashAttention-2：比1代加速1.29倍、GPU利用率从55%上升到72%，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

FlashAttention-2

    • 提出背景
      • FlashAttention-2 改进
    • 前向传播和反向传播对比
      • FlashAttention前向传播
      • FlashAttention反向传播
      • FlashAttention-2前向传播
      • FlashAttention-2反向传播
      • FlashAttention-2并行性
      • 线程束之间的工作分区
    • 总结
      • FlashAttention
      • FlashAttention-2

 

论文：https://arxiv.org/pdf/2307.08691.pdf

 

提出背景

前置：FlashAttention：高效注意力计算的新纪元

核心问题在于，如何提升 Transformer 模型在处理长序列数据时的效率，尤其是在注意力（Attention）层的计算上?

在Transformer架构中，注意力层的计算复杂度和内存需求随着输入序列长度的增加而呈二次方增长，这限制了模型的扩展性和效率。

为了解决这个问题，研究人员提出了 FlashAttention 技术，并在其基础上进一步发展出 FlashAttention-2。

 

FlashAttention-2 改进

FlashAttention-2的改进：

  1. 算法调整减少非矩阵乘法FLOPs：
    • 改进：FlashAttention-2通过减少非矩阵乘法操作，优化了计算流程，提高了利用专用硬件单元的效率。
    • 例子对比：如果FlashAttention像一个老式工厂依赖多功能但速度慢的机器，FlashAttention-2则像升级后的工厂，更多地使用专用于高速组装的机器。
    • 在FlashAttention中，非矩阵乘法操作占用了相当一部分计算资源，而在FlashAttention-2中，这些操作被优化以利用GPU的Tensor Cores，从而加快整体计算速度。

 

  1. 并行计算注意力：
    • 改进：FlashAttention-2通过在更多的线程块上并行化计算来增加GPU资源的占用率，提高了计算效率。
    • 例子对比：在FlashAttention中，注意力计算可能更集中，类似于一个分析师处理整批数据。
    • 而在FlashAttention-2中，这项工作被分成多部分，类似于多个分析师同时处理数据的不同部分，显著提高了处理速度。

 

  1. 工作分配：
    • 改进：FlashAttention-2在每个线程块内部更智能地分配工作给线程束，减少了线程束间的通信开销。
    • 例子对比：在FlashAttention中，所有的文件传递都要通过中央档案室，导致效率低下。
    • 而在FlashAttention-2中，每个部门有自己的存档点，员工可以快速获取文件，类似于在每个线程块中进行优化，使得线程束间几乎不需要通过共享内存通信，提高了效率。

这些改进的结果是FlashAttention-2的计算速度比FlashAttention快了大约2倍，达到了接近GEMM操作的效率水平。

对比原有的 FlashAttention 解法，FlashAttention-2 通过更好的工作划分和并行计算策略，显著提高了处理长序列时的效率和吞吐量，达到了更接近矩阵乘法操作的理论最大 FLOPs/s，同时减少了内存使用量，解决了处理长序列数据时面临的主要瓶颈。

 

前向传播和反向传播对比

FlashAttention前向传播

  • 数据装载：从HBM（高带宽内存）到SRAM（静态随机存取内存）加载数据块。
  • 计算：对加载的块计算注意力输出。
  • 更新输出：不写入大型中间矩阵S和P，而是直接更新输出O。
  • 优化：通过在线softmax技术进行分块计算和重缩放每个块的输出，避免了完整softmax计算，以此减少内存操作。

想象一个图书馆的书籍检索系统。

在FlashAttention的前向传播中，每当有人请求一本书时，工作人员（GPU的线程）会去一个巨大的仓库（HBM）中找到这本书，然后将它带到前台（SRAM）进行检查。

如果需要检索多本书（即处理大量数据），他们会尽量避免一次性搬运所有的书籍，而是采用分批搬运的方式来减少往返仓库的次数。

在这里插入图片描述
将键矩阵 ( K ) 和值矩阵 ( V ) 分成块，并对每个块计算注意力。

通过这种方式，计算避免了在高带宽内存（HBM）中存储完整的中间注意力矩阵 ( S ) 和 ( P )，而是在片上SRAM中进行计算，并且只存储必要的输出，这减少了内存的读/写操作。

FlashAttention反向传播

  • 重计算：重新计算注意力矩阵S和P的值，以避免在HBM中存储这些大型矩阵。
  • 内存效率：FlashAttention在反向传播中也实现了显著的内存节约和计算加速。

在图书馆的例子中，如果客户想要返回一些书籍，FlashAttention的反向传播相当于工作人员重新确认每本书的位置，然后将它们放回仓库中的正确位置。

这个过程被优化以确保最小化搬运和存储过程，以便书籍可以在下次被快速检索。

FlashAttention-2前向传播

  • 算法优化：减少了非矩阵乘法FLOPs的数量，充分利用专门的GPU计算单元（如Tensor Cores）。
  • 维护未缩放版本：在整个计算循环中维护一个未缩放的输出版本，只在最后缩放最终输出。
  • 存储优化：只存储logsumexp，而不是最大值和指数和，以简化后向传播中的计算。

在改进版的系统中，FlashAttention-2对前向传播进行了优化，就像工作人员现在使用一个电子目录（非矩阵乘法FLOPs的减少）来快速确定书籍的位置，然后只搬运所需的书籍到前台。

他们还开始记录哪些书籍最受欢迎，以此减少未来的检索时间（存储logsumexp代替max和sum）。

FlashAttention-2反向传播

  • 与前向传播类似：FlashAttention-2的反向传播与FlashAttention大致相同，但进行了一些小的调整，例如，仅使用行logsumexp L 而不是行最大值和行指数和。

在返回书籍时，工作人员现在只需要检查电子目录中的位置，然后直接将书放回正确的架子上，而不需要再次确认整个仓库的存储布局（使用行logsumexp代替完整的行最大值和行指数和计算）。

FlashAttention-2并行性

  • 序列长度并行化：在序列长度维度上进行并行化，提高了GPU资源的占用率，尤其是在批量小或头数少时。

在FlashAttention-2中，图书馆现在雇佣了更多的工作人员，并将仓库分成了几个区域。

现在，当多个请求同时到来时，不同的工作人员可以同时在不同的区域检索书籍，大大加快了整个系统的响应速度（序列长度并行化）。

在这里插入图片描述
FlashAttention-2的前向传播和反向传播的并行化策略。

在前向传播中，不同的工作单元（线程块）负责注意力矩阵的不同行块，而在反向传播中，它们处理列块。

线程束之间的工作分区

  • FlashAttention：在计算QKᵀ时将K和V分割在不同的线程束中。
  • FlashAttention-2：改进了工作分区，减少了在不同线程束之间共享内存读写的需要。

在FlashAttention中，所有工作人员在检索书籍后都需要到一个中央台账（共享内存）进行登记。

在FlashAttention-2中，这个过程被优化了。

每个工作人员都有自己的手持设备（线程束内部的工作分配），可以在找到书籍的同时进行登记，这减少了在中央台账前排队的时间。

在FlashAttention中，由于所有线程束都需要写入它们的中间结果到共享内存，同步后再加起来，这导致了前向传播中的效率损失。

而在FlashAttention-2中，通过将Q分割在不同线程束中，并保持K和V对所有线程束可访问，每个线程束只需与它们的共享片段V进行乘法运算以获得输出的一部分，从而无需线程束间通信。

在这里插入图片描述

比较了FlashAttention和FlashAttention-2在前向传播中不同线程束（一组线程）之间的工作分区。

在 FlashAttention(a) 中，线程束分割 ( K ) 和 ( V ) 同时计算 ( QK^T ) 的切片，这需要同步操作来组合结果。

FlashAttention-2(b) 通过在线程束间分割 ( Q )，同时保持所有线程束都能访问 ( K ) 和 ( V )，优化了这一点，减少了同步和共享内存访问的需求，有可能提高了计算效率。

 

总结

FlashAttention

  1. IO感知计算：FlashAttention通过优化内存访问来减少计算时间和内存占用。
  2. 分块计算（Tiling）：通过将数据分块处理，减轻了内存压力，并使注意力计算更高效。
  3. 重计算策略：反向传播时不存储大型中间矩阵S和P，而是重计算，从而节省内存。

FlashAttention-2

  1. 算法调整减少非矩阵乘法FLOPs：通过减少非矩阵乘法操作来优化计算流程，提高专用硬件单元的效率。
  2. 并行计算注意力：在更多的线程块上并行化计算来提高GPU资源的占用率，提高计算效率。
  3. 工作分配优化：在每个线程块内部更智能地分配工作给线程束，减少线程束间的通信开销。
  • FlashAttention 强调了在单个GPU上实现注意力计算的IO效率，通过分块计算和重计算来减少内存需求和提高效率。
  • FlashAttention-2 在FlashAttention的基础上进一步改进，特别是在并行计算和线程束工作分配方面做了优化，以提高在处理长序列时的效率和吞吐量。

例如，如果将FlashAttention比作一辆提升了燃油效率的汽车，那么FlashAttention-2就像是在此基础上增加了涡轮增压器，进一步提升了性能。FlashAttention-2的主要改进在于更高效地使用GPU资源，减少了内存访问和同步操作，使得模型能够在更长的序列上更快地进行训练和推理，同时还能减少模型训练和使用时的资源消耗。

这篇关于【大模型上下文长度扩展】FlashAttention-2：比1代加速1.29倍、GPU利用率从55%上升到72%的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！


http://www.chinasem.cn/article/686236。

相关文章

大模型研发全揭秘：客服工单数据标注的完整攻略

大模型研发全揭秘：客服工单数据标注的完整攻略

在人工智能（AI）领域，数据标注是模型训练过程中至关重要的一步。无论你是新手还是有经验的从业者，掌握数据标注的技术细节和常见问题的解决方案都能为你的AI项目增添不少价值。在电信运营商的客服系统中，工单数据是客户问题和解决方案的重要记录。通过对这些工单数据进行有效标注，不仅能够帮助提升客服自动化系统的智能化水平，还能优化客户服务流程，提高客户满意度。本文将详细介绍如何在电信运营商客服工单的背景下进行
阅读更多...
csu 1446 Problem J Modified LCS （扩展欧几里得算法的简单应用）

csu 1446 Problem J Modified LCS （扩展欧几里得算法的简单应用）

这是一道扩展欧几里得算法的简单应用题，这题是在湖南多校训练赛中队友ac的一道题，在比赛之后请教了队友，然后自己把它a掉 这也是自己独自做扩展欧几里得算法的题目 题意：把题意转变下就变成了：求d1*x - d2*y = f2 - f1的解,很明显用exgcd来解 下面介绍一下exgcd的一些知识点：求ax + by = c的解 一、首先求ax + by = gcd(a,b)的解 这个
阅读更多...
Andrej Karpathy最新采访：认知核心模型10亿参数就够了，AI会打破教育不公的僵局

Andrej Karpathy最新采访：认知核心模型10亿参数就够了，AI会打破教育不公的僵局

夕小瑶科技说 原创  作者 | 海野 AI圈子的红人，AI大神Andrej Karpathy，曾是OpenAI联合创始人之一，特斯拉AI总监。上一次的动态是官宣创办一家名为 Eureka Labs 的人工智能+教育公司 ，宣布将长期致力于AI原生教育。 近日，Andrej Karpathy接受了No Priors（投资博客）的采访，与硅谷知名投资人 Sara Guo 和 Elad G
阅读更多...
Retrieval-based-Voice-Conversion-WebUI模型构建指南

Retrieval-based-Voice-Conversion-WebUI模型构建指南

一、模型介绍 Retrieval-based-Voice-Conversion-WebUI（简称 RVC）模型是一个基于 VITS（Variational Inference with adversarial learning for end-to-end Text-to-Speech）的简单易用的语音转换框架。 具有以下特点 简单易用：RVC 模型通过简单易用的网页界面，使得用户无需深入了
阅读更多...
科研绘图系列：R语言扩展物种堆积图（Extended Stacked Barplot）

科研绘图系列：R语言扩展物种堆积图（Extended Stacked Barplot）

介绍 R语言的扩展物种堆积图是一种数据可视化工具，它不仅展示了物种的堆积结果，还整合了不同样本分组之间的差异性分析结果。这种图形表示方法能够直观地比较不同物种在各个分组中的显著性差异，为研究者提供了一种有效的数据解读方式。 加载R包 knitr::opts_chunk$set(warning = F, message = F)library(tidyverse)library(phyl
阅读更多...
透彻！驯服大型语言模型（LLMs）的五种方法，及具体方法选择思路

透彻！驯服大型语言模型（LLMs）的五种方法，及具体方法选择思路

引言 随着时间的发展，大型语言模型不再停留在演示阶段而是逐步面向生产系统的应用，随着人们期望的不断增加，目标也发生了巨大的变化。在短短的几个月的时间里，人们对大模型的认识已经从对其zero-shot能力感到惊讶，转变为考虑改进模型质量、提高模型可用性。 「大语言模型(LLMs)其实就是利用高容量的模型架构（例如Transformer）对海量的、多种多样的数据分布进行建模得到，它包含了大量的先验
阅读更多...
图神经网络模型介绍（1）

图神经网络模型介绍（1）

我们将图神经网络分为基于谱域的模型和基于空域的模型，并按照发展顺序详解每个类别中的重要模型。 1.1基于谱域的图神经网络         谱域上的图卷积在图学习迈向深度学习的发展历程中起到了关键的作用。本节主要介绍三个具有代表性的谱域图神经网络：谱图卷积网络、切比雪夫网络和图卷积网络。 （1）谱图卷积网络 卷积定理：函数卷积的傅里叶变换是函数傅里叶变换的乘积，即F{f*g}
阅读更多...
秋招最新大模型算法面试，熬夜都要肝完它

秋招最新大模型算法面试，熬夜都要肝完它

💥大家在面试大模型LLM这个板块的时候，不知道面试完会不会复盘、总结，做笔记的习惯，这份大模型算法岗面试八股笔记也帮助不少人拿到过offer ✨对于面试大模型算法工程师会有一定的帮助，都附有完整答案，熬夜也要看完，祝大家一臂之力 这份《大模型算法工程师面试题》已经上传CSDN，还有完整版的大模型 AI 学习资料，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费
阅读更多...
【生成模型系列（初级）】嵌入（Embedding）方程——自然语言处理的数学灵魂【通俗理解】

【生成模型系列（初级）】嵌入（Embedding）方程——自然语言处理的数学灵魂【通俗理解】

【通俗理解】嵌入（Embedding）方程——自然语言处理的数学灵魂 关键词提炼 #嵌入方程 #自然语言处理 #词向量 #机器学习 #神经网络 #向量空间模型 #Siri #Google翻译 #AlexNet 第一节：嵌入方程的类比与核心概念【尽可能通俗】 嵌入方程可以被看作是自然语言处理中的“翻译机”，它将文本中的单词或短语转换成计算机能够理解的数学形式，即向量。 正如翻译机将一种语言
阅读更多...
AI Toolkit + H100 GPU，一小时内微调最新热门文生图模型 FLUX

AI Toolkit + H100 GPU，一小时内微调最新热门文生图模型 FLUX

上个月，FLUX 席卷了互联网，这并非没有原因。他们声称优于 DALLE 3、Ideogram 和 Stable Diffusion 3 等模型，而这一点已被证明是有依据的。随着越来越多的流行图像生成工具（如 Stable Diffusion Web UI Forge 和 ComyUI）开始支持这些模型，FLUX 在 Stable Diffusion 领域的扩展将会持续下去。 自 FLU
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2