Hugging Face高性能技术五：Transformer高效推断（bitsandbytes、FlashAttention、 BetterTransformer）

  本文整合了HF文档《CPU inference》和《GPU inference》，重复部分进行了合并。另外还参考了《面向生产的 LLM 优化》等HF博客和文档。

一、 TorchScript（CPU)

TorchScript教程

  TorchScript是 PyTorch 模型（nn.Module）的一种中间表示形式，可以在高性能环境（如 C++）中运行而不需要Python环境。这有助于加速深度学习模型的推断，并提供更大的灵活性和可部署性，使其适用于在生产环境中部署和执行。以下是TorchScript的详细介绍：

1. 中间表示形式：TorchScript是一种中间表示形式，类似于抽象语法树（AST），它将PyTorch模型编码为一种静态图形式。这种表示形式是独立于Python解释器的，因此在执行时不依赖于Python的运行时环境。

2. 模型导出：要使用TorchScript，首先需要将PyTorch模型导出为TorchScript表示。您可以使用两种主要方法来实现这一点：

   • 脚本化导出：使用torch.jit.script函数，您可以将PyTorch模型的前向传播过程转换为TorchScript。这种方法将模型的Python代码转换为TorchScript表示，但有一些限制，例如不支持动态计算图。
   • 跟踪导出：使用torch.jit.trace函数，您可以跟踪模型的执行过程，并将跟踪的操作转化为TorchScript表示。这通常用于将模型的输入和输出签名固定为特定数据类型和形状。

3. ScriptModule和ScriptFunction：TorchScript导出后，您将获得一个ScriptModule对象，它是一个包含模型和其权重的TorchScript表示。ScriptModule是一个独立的可执行单位，可以在没有Python解释器的情况下运行。此外，模型中的函数可以转换为ScriptFunction，这是一个经过JIT编译的函数，可独立运行。

4. 性能提升：TorchScript的一个主要优点是性能提升。由于它是一个静态图形式，可以使用即时编译（JIT）技术进行各种优化，如操作融合、内存管理和计算图剪枝，以提高推断性能。

5. 部署和跨平台性：TorchScript可以用于在各种环境中部署模型，包括移动设备、嵌入式系统和云服务。它还支持多种前端（例如PyTorch、LibTorch、ONNX等）和多种后端（例如CPU、GPU、加速器等），从而实现更广泛的部署和跨平台兼容性。

6. 库生态系统：TorchScript生态系统包括PyTorch中的扩展，如TorchScript自定义操作和TorchScript模型库，可以帮助用户更轻松地使用TorchScript来导出和执行模型。

  关于TorchScript的详细教程可参考《Introduction to TorchScript》。在Trainer中，可以通过设置--jit_mode_eval来启用 JIT 模式的 CPU 推理。

python run_qa.py \
--model_name_or_path csarron/bert-base-uncased-squad-v1 \
--dataset_name squad \
--do_eval \
--max_seq_length 384 \
--doc_stride 128 \
--output_dir /tmp/ \
--no_cuda \
--jit_mode_eval

PyTorch >= 1.14.0，JIT 模式可以使任何模型的预测和评估受益，因为 jit.trace 中支持 dict 输入。

二、 IPEX graph optimization（Intel CPU）

  Intel Extension for PyTorch （IPEX） 在 JIT 模式下为 Intel® CPU 提供了进一步的优化，我们建议将其与 TorchScript 结合使用，以获得更快的性能。

  IPEX 图优化融合了多头注意力、Concat Linear、Linear + Add、Linear + Gelu、Add + LayerNorm 等操作。要利用这些图形优化，请先安装 IPEX：

pip install intel_extension_for_pytorch

然后在 Trainer 中设置 --use_ipex 和 --jit_mode_eval，以启用 IPEX优化的 JIT 模式：

python run_qa.py \
--model_name_or_path csarron/bert-base-uncased-squad-v1 \
--dataset_name squad \
--do_eval \
--max_seq_length 384 \
--doc_stride 128 \
--output_dir /tmp/ \
--no_cuda \
--use_ipex \
--jit_mode_eval

三、 Optimum

Optimum文档、ONNX Runtime文档

  Optimum 是 Transformer 的扩展，它提供了一组性能优化工具，用于在目标硬件上以最高效率训练和运行模型，其中就包括ONNX Runtime (ORT)。ORT 由 🤗 Optimum 支持，可以在 Transformer 中使用🤗，而无需对代码进行太多更改。

Optimum

3.1 安装

  参考Optimum Installation，Optimum的pip安装方式如下：

python -m pip install optimum

  如上图所示，optimum有6种组件用于加速。要启用这些特定的加速功能，参照下表安装依赖：（--upgrade-strategy eager 选项确保将不同的软件包升级到最新版本）

Accelerator	Installation
ONNX runtime	pip install --upgrade-strategy eager optimum[onnxruntime]
Intel Neural Compressor (INC)	pip install --upgrade-strategy eager optimum[neural-compressor]
Intel OpenVINO	pip install --upgrade-strategy eager optimum[openvino,nncf]
Habana Gaudi Processor (HPU)	pip install --upgrade-strategy eager optimum[habana]
FuriosaAI	pip install --upgrade-strategy eager optimum[furiosa]

3.2 CPU推理

  ORT 是在CPU推理时默认的模型加速器，要使用ORT，您只需将🤗 Transformers的AutoClass替换为对应任务的ORTModel，并加载以ONNX格式保存的checkpoint。例如对问答任务运行推理时，请加载包含 model.onnx 文件的 optimum/roberta-base-squad2 checkpoint 。

from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForQuestionAnswering

model = ORTModelForQuestionAnswering.from_pretrained("optimum/roberta-base-squad2")
tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")

onnx_qa = pipeline("question-answering", model=model, tokenizer=tokenizer)

question = "What's my name?"
context = "My name is Philipp and I live in Nuremberg."
pred = onnx_qa(question, context)

  上面所述只是一个简短的示例，关于 ORT推理的更多信息，请参考指南《 Optimum Inference with ONNX Runtime》。

  另外，如果您有 Intel CPU，请查看 🤗 Optimum Intel，它支持各种压缩技术（量化、剪枝、知识蒸馏）以及将模型转换为 OpenVINO 格式以实现更高性能的推理。

3.3 GPU推理

  ORT也支持在 Nvidia GPU 上加速推理。其优化技术包括将常见操作融合到单个节点和常量折叠中（fusing common operations into a single node and constant folding），以减少执行的计算次数并加快推理速度。ORT 还会将计算密集型操作放在 GPU 上，其余操作放在 CPU 上，以便在两个设备之间智能地分配工作负载。

  要使用ORT，您只需将🤗 Transformers的AutoClass替换为对应任务的ORTModel，并指定provider参数（CUDAExecutionProvider或 TensorrtExecutionProvider）。如果您想加载一个尚未导出为ONNX格式的模型，可以设置export=True，以便即时将模型转换为ONNX格式：

from optimum.onnxruntime import ORTModelForSequenceClassification

ort_model = ORTModelForSequenceClassification.from_pretrained(
  "distilbert-base-uncased-finetuned-sst-2-english",
  export=True,
  provider="CUDAExecutionProvider",
)

现在，可以自由使用该模型进行推理：

from optimum.pipelines import pipeline
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

pipeline = pipeline(task="text-classification", model=ort_model, tokenizer=tokenizer, device="cuda:0")
result = pipeline("Both the music and visual were astounding, not to mention the actors performance.")

四、模型量化（GPU）

• 4.1 至4.3 节参考《大规模 Transformer 模型 8 比特矩阵乘简介 - 基于 Hugging Face Transformers、Accelerate 以及 bitsandbytes》
• 4.4 节参考HF文档量化文档《Quantize 🤗 Transformers models》
• 4.5 节参考《面向生产的 LLM 优化》
• 4.6 节见《Accelerate 0.24.0文档 三：超大模型推理（Sharded checkpoints、deepspeed、bitsandbytes量化）》
• 4.7 节参考量化文档《Quantize 🤗 Transformers models》、《大语言模型的模型量化(INT8/INT4)技术》

  本文只介绍量化模型的推理，且只介绍bitsandbytes库实现的LLM.int8()量化技术。

  量化后的transformer模型不支持直接训练，但是你可以使用参数高效微调方法（PEFT）微调量化模型，训练adapters。例如colab notebook《bnb-4bit-training.ipynb》中，演示了如何微调4bit模型，更多信息详见peft 。

4.1 背景

  语言模型一直在变大。例如PaLM 有 5400 亿参数，OPT、GPT-3 和 BLOOM 有大约 1760 亿参数，而且我们仍在继续朝着更大的模型发展。下图总结了最近的一些语言模型的尺寸。

  迄今为止，市面上显存最大的 GPU 芯片是 80GB 显存的 A100，因此这些模型无法在单个设备上运行。举个例子，如果我们使用 BLOOM-176B 模型的 Bfloat16 版本，其大小为$176\times 10^9\times 2byte=352GB$。仅推理 BLOOM-176B 模型，就需要 8 个 80GB A100 。而如果要微调的话，则需要更多。

  在《Hugging Face高效训练技术一：单 GPU 高效训练》第四章混合精度训练中，有介绍TF32、FP32、FP16、BF16等浮点数类型（见上图），以及为了加速训练而使用的混合精度训练技术。训练时我们使用FP16的weights, activations,gradients进行前后向传播及梯度计算，而为了避免精度溢出和舍入误差，还需要备份FP32权重用于权重更新。而对于推理，半精度(FP16)权重就能达到与 FP32 相似的精度，这样我们仅需一半 GPU 显存就能获得相同的结果。

  量化对于文本生成特别有效，因为我们关心的是选择最可能的下一个词元的分布 ，而不真正关心下一个词元的确切 logit 值。所以，只要下一个词元 logit 大小顺序保持相同， argmax 或 topk 操作的结果就会相同。

  既然使用半精度进行推理，可以缩减一半的显存，那么我们自然会想到采用8位量化，将浮点值映射为更紧凑的 Int8类型（1字节，取值范围[-127,127]），此时内存会减少到原来的1/4。但是这次会出现问题，因为直接丢弃另一半位宽来进行8位量化，推理结果的质量会急剧下降。

4.2 基础量化技术

  量化过程是从一种数据类型“舍入”到另一种数据类型，是一种有损压缩，会导致信息丢失。两种最常见的 8 位量化技术是零点量化 (zero-point quantization) 和最大绝对值 (absolute maximum quantization，absmax) 量化，其第一步都是用量化常数对输入进行归一化缩放。

1. 零点量化
   • 量化：假设原数值范围是 [-1.0,1.0]，现在需要量化到[-127,127]，那么需要先缩放 127倍，然后四舍五入到 8 位精度。比如0.3缩放为 0.3*127 = 38.1，四舍五入后得到值 38。
   • 恢复：将量化后的值反向缩放127倍，即38/127=0.2992。可见最终会有 0.008 的量化误差。这些看似微小的误差在沿着模型各层传播时往往会累积和增长，从而导致最终的精度下降。

   零点量化一大特点是零点调整，即通过缩放将原始的数值范围映射为量化后的数值范围之外，还需要通过平移将映射后的数据的最小值对齐为目标值域的最小值，这一步称为零点调整，上述例子没有体现出来。

2. absmax 量化
   • 量化：先除以张量的最大绝对值，然后再乘以数据类型的最大可表示值。
     假设你要用 absmax 对向量 [1.2, -0.5, -4.3, 1.2, -3.1, 0.8, 2.4, 5.4] 进行量化。首先需要计算该向量元素的最大绝对值，在本例中为 5.4。 Int8 的范围为 [-127, 127]，因此我们将 127 除以 5.4，得到缩放因子 23.5。最后，将原始向量乘以缩放因子得到最终的量化向量 [28, -12, -101, 28, -73, 19, 56, 127]。
   • 恢复：将 int8 量化值除以缩放因子，但由于上面的过程是“四舍五入”的，我们将丢失一些精度。
     
3. 矩阵量化技巧
   当进行矩阵乘法时，我们可以通过组合各种技巧，例如逐行或逐向量量化，来获取更精确的结果。例如对矩阵乘法 $A\times B=C$，我们不会直接使用常规量化方式，即用整个张量的最大绝对值对张量进行归一化，而会转而使用向量量化方法，找到 A 的每一行和 B 的每一列的最大绝对值，然后逐行或逐列归一化 A 和 B 。最后将 A 与 B 相乘得到 C。最后，我们再计算与 A 和 B 的最大绝对值向量的外积，并将此与 C 求哈达玛积来反量化回 FP16。有关此技术的更多详细信息可以参考 LLM.int8() 论文 或 Tim 的博客 关于量化和涌现特征的博文。

  虽然这些基本技术能够帮助我们量化深度学习模型，但它们通常会导致大模型准确性的下降。集成在Transformers 和 Accelerate中的 LLM.int8() 是第一个适用于大模型 (如 BLOOM-176B) 且不会降低准确性的量化技术。

4.3 LLM.int8()：大语言模型的零退化矩阵乘法

LLM.int8()论文《8-bit Matrix Multiplication for Transformers at Scale》、

4.3.1 算法

  超出某个分布范围的值通常称为离群值，虽然离群值特征也存在于较小的模型中，但在大于 6B 的 transformer 模型中，我们观察到几乎每层都会出现超出特定阈值的离群点，这些离群点呈现出一定的系统性模式。 直接对模型进行8位量化而导致性能下降的情况，正是由离群特征 (outlier feature) 引起的。

  8 位精度的动态范围极其有限，因此量化具有多个大值的向量会产生严重误差。此外，由于 transformer 架构的固有特性，它会将所有元素互相关联起来，这样的话，这些误差在传播几层后往往会混杂在一起。对此，我们发明了混合精度分解的方法，以对此类极端离群值进行有效量化，整个算法过程如下：

1. 从输入的隐含状态中，按列提取异常值 (使用自定义阈值提取离群值)，将矩阵分解为两部分。
2. 离群值部分使用 FP16 表示，因此它是一个经典的矩阵乘法。其余值使用int8表示，所以是8 位矩阵乘法。

   如上一节矩阵量化技巧所示，8 位矩阵乘法是通过使用向量量化将权重和隐含状态分别量化为 8 位精度，即按行量化权重矩阵，并按列量化隐含状态，然后再进行相应向量乘加操作。

3. 反量化8 位矩阵乘法结果至半精度，以便与离群值矩阵乘法的结果相加，获得最终的 FP16 结果。
   

4.3.2 实验

我们使用 lm-eval-harness 在 8 位和原始模型上运行了几个常见的基准测试，结果如下：

1. OPT-175B 模型

测试基准名	指标	指标值 - int8	指标值 - fp16	标准差 - fp16	差值
hellaswag	acc_norm	0.7849	0.7849	0.0041	0
hellaswag	acc	0.5921	0.5931	0.0049	0.001
piqa	acc	0.7965	0.7959	0.0094	0.0006
piqa	acc_norm	0.8101	0.8107	0.0091	0.0006
lambada	ppl	3.0142	3.0152	0.0552	0.001
lambada	acc	0.7464	0.7466	0.0061	0.0002
winogrande	acc	0.7174	0.7245	0.0125	0.0071

2. BLOOM-176 模型

测试基准名	指标	指标值 - int8	指标值 - fp16	标准差 - fp16	-
hellaswag	acc_norm	0.7274	0.7303	0.0044	0.0029
hellaswag	acc	0.5563	0.5584	0.005	0.0021
piqa	acc	0.7835	0.7884	0.0095	0.0049
piqa	acc_norm	0.7922	0.7911	0.0095	0.0011
lambada	ppl	3.9191	3.931	0.0846	0.0119
lambada	acc	0.6808	0.6718	0.0065	0.009
winogrande	acc	0.7048	0.7048	0.0128	0

  可以看到上述这些模型的性能下降为 0（指标的绝对差异均低于原始模型的标准误差），所以称之为0 退化。更多性能对比实验，请查看 LLM.int8() 论文 。

  实验发现使用了 LLM.int8() 的 BLOOM-176B 比 FP16 版本慢了大约 15% 到 23%（见下表），但LLM.int8() 方法的主要目的是在不降低性能的情况下降低大模型的应用门槛（减少内存），所以这个结果是可以接受的。

精度	模型	参数量	硬件	每词元延迟 (单位: 毫秒，batch size: 1)	每词元延迟 (单位: 毫秒，batch size: 8)	每词元延迟 (单位: 毫秒，batch size: 32)
bf16	BLOOM-176B	176B	8xA100 80GB	239	32	9.9
int8	BLOOM-176B	176B	4xA100 80GB	282	37.5	10.2
bf16	BLOOM-176B	176B	14xA100 40GB	285	36.5	10.4
int8	BLOOM-176B	176B	5xA100 40GB	367	46.4	oom
fp16	T5-11B	11B	2xT4 15GB	11.7	1.7	0.5
int8	T5-11B	11B	1xT4 15GB	43.5	5.3	1.3
fp32	T5-3B	3B	2xT4 15GB	45	7.2	3.1
int8	T5-3B	3B	1xT4 15GB	312	39.1	10.2

有关 8 位量化基础概念的更多信息，请参考博客《Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using Hugging Face Transformers, Accelerate and bitsandbytes》

4.4 bitsandbytes量化（transformer 4.35.0）

参考《Quantize 🤗 Transformers models》

4.4.1 一般用法

1. 使用 from_pretrained() 加载量化模型

  您可以在调用 from_pretrained() 方法时使用 load_in_8bit 或 load_in_4bit 参数来量化模型，只要您的模型支持使用 🤗 Accelerate 加载并包含 torch.nn.Linear 层。然后，像通常使用 PreTrainedModel 一样使用您的模型。

  Transformers 开箱即用地支持简单的流水线并行。为此，只需使用 device_map="auto"加载模型，它就会自动将不同层放到相应的 GPU 上，使模型的加载和推理更加灵活和高效。(目前 device_map = 'auto’只适合推理）

from transformers import AutoModelForCausalLM

model_8bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m",device_map="auto", load_in_8bit=True)
model_4bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", device_map="auto",load_in_4bit=True)

  另外，默认情况下，全连接层之外的所有其他模块（例如 torch.nn.LayerNorm ）都将被转为torch.float16格式，不过你也可以覆盖 torch_dtype 参数改为其它的数据类型。

from transformers import AutoModelForCausalLM

model_name = "bigscience/bloom-2b5"
model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True,torch_dtype=torch.float32)
model_8bit.model.decoder.layers[-1].final_layer_norm.weight.dtype 

torch.float32

您可以使用 get_memory_footprint 方法检查模型的内存占用情况。

print(model.get_memory_footprint())

2. 在 🤗 Hub 上推送量化模型（bitsandbytes>0.37.2 ）
   您可以简单地使用 push_to_hub 方法将量化模型推送到 Hub 上。这将首先推送量化配置文件，然后推送量化模型权重。

   from transformers import AutoModelForCausalLM, AutoTokenizer
   
   model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", device_map="auto", load_in_8bit=True)
   tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
   
   model.push_to_hub("bloom-560m-8bit")
   

3. 从 🤗 Hub 加载量化模型
   您可以使用 from_pretrained 方法从 Hub 加载量化模型，加载前确保模型配置对象中含有quantization_config属性。这种情况下，无需指定参数 load_in_8bit=True，但需要确保安装了 bitsandbytes 和 accelerate 。

   from transformers import AutoModelForCausalLM, AutoTokenizer
   
   model = AutoModelForCausalLM.from_pretrained("{your_username}/bloom-560m-8bit", device_map="auto")
   

4.4.2 量化模型高级用例（4-bit）

下面将介绍一些可以使用 FP4 量化执行的高级用例

1. 更改compute dtype
   compute dtype 用于更改将在计算期间使用的 dtype，默认情况下为 float32，我们可以设置为 bf16 以实现加速：

   import torch
   from transformers import BitsAndBytesConfig
   
   quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
   

2. 使用 NF4（普通浮点 4）数据类型
   NF4 数据类型是一种新的 4 位数据类型，适用于已使用正态分布初始化的权重。

   from transformers import BitsAndBytesConfig
   
   nf4_config = BitsAndBytesConfig(
       load_in_4bit=True,
       bnb_4bit_quant_type="nf4",
   )
   
   model_nf4 = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=nf4_config)
   

3. 使用嵌套量化实现更高效的内存推理
   嵌套量化技术可以在不增加性能的情况下节省更多内存 - 根据我们的经验观察，这可以在序列长度为 1024、批量大小为 1、梯度累积步长为 4 的 NVIDIA-T4 16GB 上微调 llama-13b 模型。

   from transformers import BitsAndBytesConfig
   
   double_quant_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_use_double_quant=True,
   )
   
   model_double_quant = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=double_quant_config)
   

4.4.3 量化模型高级用例（8-bit）

1. 卸载到CPU
   对于大模型来说，如果 GPU 上没有足够的空间将整个模型存储在 GPU 上，则可以将权重分配到CPU和GPU之间。要注意的是，分配到CPU上的权重不会被转换为8位，而是保持在float32格式。

  首先，从transformers库中加载BitsAndBytesConfig，并将属性llm_int8_enable_fp32_cpu_offload设置为True。这个配置对象的目的是启用将模型权重分配到CPU上时，保持在float32格式而不是转换为8位。

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)

  假设你想加载bigscience/bloom-1b7模型，并且你的GPU内存足够容纳整个模型，除了lm_head部分。此时，需要编写一个自定义的device_map，以便明确指定权重在哪些设备上加载。

device_map = {
    "transformer.word_embeddings": 0,
    "transformer.word_embeddings_layernorm": 0,
    "lm_head": "cpu",
    "transformer.h": 0,
    "transformer.ln_f": 0,
}

加载模型：

model_8bit = AutoModelForCausalLM.from_pretrained(
    "bigscience/bloom-1b7",
    device_map=device_map,
    quantization_config=quantization_config,
)

2. 更改异常值阈值

  “异常值”指的是隐藏状态值大于某个特定阈值的值，这对应于LLM.int8()论文中描述的异常值检测中的异常值阈值。

  隐藏状态值通常呈正态分布，即大多数值在范围[-3.5, 3.5]内，但对于大型模型，存在一些特殊的系统性异常值，其分布差异很大，这些异常值通常在区间[-60, -6]或[6, 60]内，对这些异常值的操作将在fp16下执行。

   Int8 量化对于幅度约为 5 的值效果很好，但超过这个范围，会出现显著的性能损失。建议的默认阈值是6，但对于更不稳定的模型（例如小模型、微调模型），可能需要使用较低的阈值。该参数可以影响模型的推理速度，建议根据实际用例调整此参数以找到最佳值。

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "bigscience/bloom-1b7"

quantization_config = BitsAndBytesConfig(
    llm_int8_threshold=10,
)

model_8bit = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map=device_map,
    quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

  关于异常值特殊处理的原理，详见LLM.int8()论文，或者博客《Hugging Face高性能技术五：Transformer高效推断（bitsandbytes、FlashAttention、 BetterTransformer）》。

3. 跳过某些模块
   模型中有些模块为了保持稳定性，不能转换为8位。例如，Jukebox模型有几个lm_head模块，在量化时应该被跳过，这可以通过调整llm_int8_skip_modules参数进行设置。

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "bigscience/bloom-1b7"

quantization_config = BitsAndBytesConfig(
    llm_int8_skip_modules=["lm_head"],
)

model_8bit = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map=device_map,
    quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

4.4.4 微调量化模型

   8 位或 4 位量化模型无法执行全量训练。但是，您可以利用参数高效微调方法 (PEFT) 来微调这些模型，详见peft Github示例：

• Colab notebook《Finetune-opt-bnb-peft.ipynb》：使用peft库的LoRa功能微调量化后的OPT-6.7b模型
• 《Finetuning Whisper-large-V2 on Colab using PEFT-Lora + BNB INT8 training + Streaming dataset》
• 《Finetuning Whisper-large-V2 on Colab using PEFT-Lora + BNB INT8 training》

  请注意，device_map=auto 仅用于推理。这是因为推理过程通常不需要进行梯度计算，而且模型参数已经在训练期间被优化，因此在推理时可以更灵活地选择设备。

  加载模型进行训练时，不需要显式传递device_map参数。系统会自动将模型加载到GPU上进行训练。如果需要，您可以将设备映射设置为特定设备，例如cuda:0, 0, torch.device('cuda:0')。

4.5 bitsandbytes量化测试

参考《面向生产的 LLM 优化》

4.5.1 安装依赖

   LLM.int8() 方法集成在bitsandbytes 库的Linear8bitLt 模块中（torch.nn.modules 的子类），使用前，需要先进行安装。

# 此版本支持 8-bit 和 4-bit
pip install bitsandbytes>=0.39.0 accelerate>=0.20.0

# install Transformers
pip install transformers

4.5.2 测试bfloat16模型（octocoder）

  如果你不确定 Hub 上的模型权重的精度如何，可随时查看模型配置文件内的 torch_dtype 项。下面，我们测试 bigcode/octocoder 模型，因为它可以在单个 40GB A100 GPU 上运行。首先测试bfloat16 精度，根据上面的速算公式，预计推理所需的显存约为 31 GB。

  from_pretrained(..., torch_dtype=...) 接口默认精度为 float32，但现在几乎所有模型都是用 bfloat16 训练的，如果 你的 GPU 支持 bfloat16 的话，你就不应该以 float32 来运行推理。float32 并不会提供比训练精度更好的推理结果，而且还会占用更多的内存。

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

4.5.3 定义辅助函数

1. 定义bytes_to_giga_bytes函数，查看显存占用：

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

调用 torch.cuda.max_memory_allocated 来测量 GPU 显存的峰值占用：

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

29.0260648727417

这个结果与我们预计的比较接近，因为从字节到千字节需要乘以 1024 而不是 1000。因此，速算公式也可以理解为“最多 X GB”。

2. 定义 flush函数，释放所有已分配的显存，便于后续测试。

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.cuda.empty_cache()
  torch.cuda.reset_peak_memory_stats()

在最新的 accelerate 库中，你还可以使用名为 release_memory() 的方法。

from accelerate.utils import release_memory
# ...

release_memory(model)

4.5.4 测试8位模型

下面我们加载8位模型，再次运行我们的示例，并测量其显存使用情况。

flush()

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
resul

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

可以看到结果和之前的一样，说明准确性没有损失！我们看一下这次用了多少显存。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated()

15.219234466552734

显存明显减少！降至 15GB 多一点，这样就可以在 4090 这样的消费级 GPU 上运行该模型了。同时，我们也注意到推理速度出现了些许减慢。

from transformers import AutoModelForCausalLM

model_name = "bigscience/bloom-2b5"
model_4bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_4bit=True)

4.5.5 测试4位模型

del model
del pipe

flush()

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True, low_cpu_mem_usage=True, pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

输出几乎与以前相同 - 只是在代码片段之前缺了 python 这个词。我们看下需要多少显存。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

9.543574333190918

  仅需 9.5GB！这样模型就可以在 RTX3090、V100 和 T4 等大多数人都可以轻松获取的 GPU 上运行。

  虽然我们这里看到模型的准确性几乎没有下降，但实际上，4位模型通常相比8位模型或者bfloat16模型，推理结果会有所变化，是否使用，看用户的选择。另外，由于 4 比特量化使用了更激进的量化方法，所以其推理速度会比8位模型更慢。

  最后，GPTQ量化方法，甚至可以将模型量化为 3 或 2 比特，对输出的影响仍可接受，详见HF量化文档《Quantize 🤗 Transformers models》，这里面也包含了量化的更多信息。总之，量化方案旨在降低内存，同时尽量保持模型的推理结果尽可能准确 ( 即尽可能接近 bfloat16)。

4.5.6 量化自定义模型

  bitsandbytes还支持将任何精度的 checkpoint 或模型转换为 8 位 ，但目前，仅当模型的输入张量数据类型为 FP16 时， Int8 模块才能工作。下面演示使用 bitsandbytes 将一个小模型转换为 int8 。

import torch
import torch.nn as nn

import bitsandbytes as bnb
from bnb.nn import Linear8bitLt

fp16_model = nn.Sequential(
    nn.Linear(64, 64),
    nn.Linear(64, 64)
)
# 假设已经在你的数据集和任务上训完了你的模型，现在需要进行保存
[... train the model ...]
torch.save(fp16_model.state_dict(), "model.pt")

先定义一个 int8 模型：

int8_model = nn.Sequential(
    Linear8bitLt(64, 64, has_fp16_weights=False),
    Linear8bitLt(64, 64, has_fp16_weights=False)
)

int8_model.load_state_dict(torch.load("model.pt"))
int8_model[0].weight

可以看到此时还是fp16模型权重：

Parameter containing:
tensor([[ 0.0031, -0.0438, 0.0494, ..., -0.0046, -0.0410, 0.0436],
        [-0.1013, 0.0394, 0.0787, ..., 0.0986, 0.0595, 0.0162],
        [-0.0859, -0.1227, -0.1209, ..., 0.1158, 0.0186, -0.0530],
        ...,
        [ 0.0804, 0.0725, 0.0638, ..., -0.0487, -0.0524, -0.1076],
        [-0.0200, -0.0406, 0.0663, ..., 0.0123, 0.0551, -0.0121],
        [-0.0041, 0.0865, -0.0013, ..., -0.0427, -0.0764, 0.1189]],
       dtype=torch.float16)

  默认情况下，has_fp16_weights=True，用于在训练时能进行 Int8/FP16 混合精度。但是，在推理中，我们需要节省内存，因此需要设置为False。

  一旦将模型的设备设置为 GPU，调用 .to函数后，量化过程已完成：

# 量化模型
int8_model = int8_model.to(torch.device('cuda', 0)) 
int8_model[0].weight

Parameter containing:
tensor([[ 3, -47, 54, ..., -5, -44, 47],
        [-104, 40, 81, ..., 101, 61, 17],
        [ -89, -127, -125, ..., 120, 19, -55],
        ...,
        [ 82, 74, 65, ..., -49, -53, -109],
        [ -21, -42, 68, ..., 13, 57, -12],
        [ -4, 88, -1, ..., -43, -78, 121]],
        device='cuda:0', dtype=torch.int8, requires_grad=True)

现在你只需将输入推给正确的 GPU 并确保输入数据类型是 FP16 的，你就可以使用该模型进行推理了:

input_ = torch.randn(64, dtype=torch.float16)
hidden_states = int8_model(input_.to(0))

比较 fp16 模型与 int8 模型输出结果：

hidden_states_fp16 = fp16_model(input_.to(0))
print(torch.max(hidden_states_fp16 - hidden_states_int8))

最大的绝对误差为

tensor(0.0098, device='cuda:0', dtype=torch.float16, grad_fn=<MaxBackward1>)

  T5-11B 模型的 checkpoint 精度为 FP32，需要 42GB 内存，Google Colab 里跑不动。使用我们的 8 位模块，它仅需 11GB 内存，因此能轻易跑通，完整代码见T5-11B Colab demo。

4.6 bitsandbytes量化（accelerate 0.24.0）

  Accelerate库也集成了bitsandbytes 量化功能，几行代码就可以实现4位量化和8位量化。详见《Accelerate 0.24.0文档 三：超大模型推理（内存估算、Sharded checkpoints、deepspeed、bitsandbytes量化、分布式推理）》

4.7 其它量化方法（有空再补）

  除了上面讲解的之外，还有SmoothQuant、GPTQ 、AWQ等量化方法，详见HF量化文档《Quantize 🤗 Transformers models》、知乎《大语言模型的模型量化(INT8/INT4)技术》。

五、 BetterTransformer

参考BetterTransformer文档、《面向生产的 LLM 优化》

5.1 BetterTransformer简介

  BetterTransformer通过其 fastpath（Transformer 函数的原生 PyTorch 实现）来加速推断。fastpath的两个优化包括：

1. 融合（fusion）：将多个顺序操作合并成一个kernel以减少计算步骤的数量。
2. 避免 padding tokens固有的稀疏性，以避免使用嵌套张量进行不必要的计算。

  BetterTransformer还将所有注意力操作转换为更节省内存的缩放点积注意力（ SDPA， scaled dot product attention），更多介绍详见《A BetterTransformer for Fast Transformer Inference》。

并非所有模型都支持 BetterTransformer。查看此列表以查看模型是否支持 BetterTransformer。

  如本文第三章所示， Optimum已经集成了BetterTransformer，且支持CPU和GPU推理。首先安装 Optimum，然后使用BetterTransformer.transform()将模型转为BetterTransformer：

from transformers import AutoModelForCausalLM
from optimum.bettertransformer import BetterTransformer

model = AutoModelForCausalLM.from_pretrained("roberta-base")
# 或尝试以accelerate加载模型
# model = AutoModel.from_pretrained("roberta-base", device_map="auto")
model = BetterTransformer.transform(model)

默认情况下，这将覆盖你的模型。如果您出于某些原因想保留它，只需添加标志 keep_original_model=True 即可！

from optimum.bettertransformer import BetterTransformer

model_bt = BetterTransformer.transform(model, keep_original_model=True)

Transformer’s pipeline 也与此集成兼容，您可以将其用作 BetterTransformer 流水线的加速器：

# CPU推理
from optimum.pipelines import pipeline

pipe = pipeline("fill-mask", "distilbert-base-uncased", accelerator="bettertransformer")
pipe("I am a student at [MASK] University.")

如果要在 GPU 设备上运行管道，请运行：

# GPU推理
from optimum.pipelines import pipeline

pipe = pipeline("fill-mask", "distilbert-base-uncased", accelerator="bettertransformer", device=0)
...

也可以像往常一样使用 transformers.pipeline ，直接传递转换后的BetterTransformer模型：

from transformers import pipeline

pipe = pipeline("fill-mask", model=model_bt, tokenizer=tokenizer, device=0)
...

  训练兼容性：如果要将其还原为原始的Transformers 模型，可以使用 reverse_bettertransformer() 方法，再将其以标准Transformers 模型进行保存。

model = model.reverse_bettertransformer()
model.save_pretrained("saved_model")

5.2 BetterTransformer测试（OctoCoder）

  SDPA（缩放点积注意力）还可以在后台调用FlashAttention内核。由于FlashAttention显著减少对较慢的高带宽显存的需求，而更多使用了更快的片上内存 (SRAM)，所以与默认注意力相比，Flash 注意力的推理速度要快得多。有关FlashAttention的详细信息，会在下一章进行介绍。

  要启用 FlashAttention 或检查它在给定设置（硬件、问题大小）中是否可用，请使用torch.backends.cuda.sdp_kernel 上下文管理器来进行查看。

  下面看一个实际的例子。继续以4.5节中的 OctoCoder模型为例，我们输入更长的提示进行注意力测试，其中包括所谓的“系统提示”（引导 LLM 去适应特定的用户任务）。下面的例子中，系统提示引导 OctoCoder 成为更好的编程助手。

system_prompt = """Below are a series of dialogues between various people and an AI technical assistant.
The assistant tries to be helpful, polite, honest, sophisticated, emotionally aware, and humble but knowledgeable.
The assistant is happy to help with code questions and will do their best to understand exactly what is needed.
It also tries to avoid giving false or misleading information, and it caveats when it isn't entirely sure about the right answer.
That said, the assistant is practical really does its best, and doesn't let caution get too much in the way of being useful.

The Starcoder models are a series of 15.5B parameter models trained on 80+ programming languages from The Stack (v1.2) (excluding opt-out requests).
The model uses Multi Query Attention, was trained using the Fill-in-the-Middle objective, and with 8,192 tokens context window for a trillion tokens of heavily deduplicated data.

-----

Question: Write a function that takes two lists and returns a list that has alternating elements from each input list.

Answer: Sure. Here is a function that does that.

def alternating(list1, list2):
   results = []
   for i in range(len(list1)):
       results.append(list1[i])
       results.append(list2[i])
   return results

Question: Can you write some test cases for this function?

Answer: Sure, here are some tests.

assert alternating([10, 20, 30], [1, 2, 3]) == [10, 1, 20, 2, 30, 3]
assert alternating([True, False], [4, 5]) == [True, 4, False, 5]
assert alternating([], []) == []

Question: Modify the function so that it returns all input elements when the lists have uneven length. The elements from the longer list should be at the end.

Answer: Here is the modified function.

def alternating(list1, list2):
   results = []
   for i in range(min(len(list1), len(list2))):
       results.append(list1[i])
       results.append(list2[i])
   if len(list1) > len(list2):
       results.extend(list1[i+1:])
   else:
       results.extend(list2[i+1:])
   return results

-----
"""

  为了演示需要，我们将系统提示复制十倍，使输入长度足够长以观察 Flash 注意力带来的内存节省。然后再加上原始提示 "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here" :

long_prompt = 10 * system_prompt + prompt

以 bfloat16 精度再次初始化模型。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

运行模型，同时测量其峰值 GPU 显存及推理时间。

import time

start_time = time.time()
result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

Generated in 10.96854019165039 seconds.
Sure. Here is a function that does that.\n\ndef bytes_to_giga(bytes):\n return bytes / 1024 / 1024 / 1024\n\nAnswer: Sure. Here is a function that does that.\n\ndef

  输出与之前一样，但是这一次，模型会多次重复答案（系统提示重复了十次），直到达到 60 个词元为止。我们测量一下峰值 GPU 显存需求。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

37.668193340301514

  正如我们所看到的，因为输入序列变长了，峰值 GPU 显存需求现在明显高于以前，整个生成过程也需要一分多钟的时间。下面调用 flush() 来释放 GPU 内存以便测试BetterTransformers。

flush()

  我们也可以使用PreTrainedModel.to_bettertransformer() 方法将模型转换为 BetterTransformers，这会因此而启用 PyTorch 的 SDPA 自注意力，其实现正是基于 Flash 注意力的。为便于比较，我们运行相同的函数。

model.to_bettertransformer()

# 现在我们运行与之前完全相同的代码片段，但此时 Transformers 在底层将使用 Flash 注意力。
start_time = time.time()
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    result = pipe(long_prompt, max_new_tokens=60)[0]["generated_text"][len(long_prompt):]

print(f"Generated in {time.time() - start_time} seconds.")
result

Generated in 3.0211617946624756 seconds.
 Sure. Here is a function that does that.\n\ndef bytes_to_giga(bytes):\n return bytes / 1024 / 1024 / 1024\n\nAnswer: Sure. Here is a function that does that.\n\ndef

  结果与之前完全相同，但由于 Flash 注意力，我们可以观察到非常显著的加速。接着测量一下内存消耗。

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

32.617331981658936

  我们可以观察到，与刚开始的短输入序列相比，使用 Flash 注意力且输入长序列时，我们只多用了大约 100MB 的 GPU 显存。

六、FlashAttention-2（GPU)

• 论文《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》、FlashAttention github地址
• 博客《面向生产的 LLM 优化》、《从 FlashAttention 到 PagedAttention, 如何进一步优化 Attention 性能》

6.1 FlashAttention：速度飞跃

6.1.1 背景

  当今表现最好的 LLM 其基本架构大体相似，包括前馈层、激活层、层归一化层以及最重要的自注意力层。其中，自注意力层是大语言模型 (LLM) 的核心，因为其使模型能够理解输入词元之间的上下文关系。然而，自注意力层在计算以及峰值显存这两个方面都随着输入词元的数目 (也称为序列长度，下文用 N 表示) 呈二次方增长，这对于较长的输入序列 (如16000 个输入词元) 来说，会成为一个严重的问题。

  我们仔细分析一下。计算长度为 N 的输入序列$\mathbf{X}$ 的自注意力层的输出 $\mathbf{O}$ ，其公式为:

$$Attention(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

上述公式可拆解为：

$$\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }\in {\mathbb{R}}^{N\times N},\mathbf{P}=\mathrm{softmax}(\mathbf{S})\in {\mathbb{R}}^{N\times N},\mathbf{O}=\mathbf{PV}\in {\mathbb{R}}^{N\times d}$$
  其中，$\mathbf{Q}={\mathbf{W}}_q\mathbf{X},\mathbf{V}={\mathbf{W}}_v\mathbf{X},\mathbf{K}={\mathbf{W}}_k\mathbf{X}$，$\mathbf{X}=({\mathbf{x}}_1,\dots {\mathbf{x}}_N)$ 是注意力层的输入序列，所以$\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$（ $N$表示序列长度，$d$ 表示head维度），所以乘积 $\mathbf{S}$是$O(N^2)$的操作。另外还有几个带宽约束的操作：

  标准的attention实现会将矩阵$\mathbf{S}$和$\mathbf{P}$实体化到HBM中,这需要$O(N^2)$的内存，通常N >> d(例如,对于GPT2,N=1024,d=64)。大多数运算是受内存限制的，比如逐元素运算的mask 和 softmax 操作，对 $\mathbf{P}$ 的 dropout 操作。下图算法展示了 HBM 与 SRAM 之间的数据传输过程：

  LLM 通常有多个注意力头，因此可以并行进行多个自注意力计算。 假设 LLM 有 40 个注意力头并以 bfloat16 精度运行，我们可以计算出存储 $\mathbf{Q}{\mathbf{K}}^{\mathbf{T}}$矩阵的内存需求为 $40\times 2\times N^2$字节。当 $N=1000$ 时仅需要大约 50MB 的显存，但当 $N=16000$ 时，我们需要 19GB 的显存，当$N=100,000$ 时，仅存储$\mathbf{Q}{\mathbf{K}}^{\mathbf{T}}$矩阵就需要近 1TB。可见，随着输入上下文越来越长，默认的自注意力算法所需的内存很快就会变得非常昂贵。

  我们如何摆脱长输入文本对内存的过高要求？ Tri Dao 等人 开发了这样一种新算法，并将其称为 Flash 注意力。

  简而言之，Flash 注意力将$\mathbf{V}\times \text{Softmax}({\mathbf{QK}}^T)$的计算分解成若干步骤，通过迭代多个 softmax 计算步来将输出分成多个较小的块进行计算:

$${\text{O}}_i\leftarrow s_{ij}^a\times {\text{O}}_i+s_{ij}^b\times {\mathbf{V}}_j\times \text{Softmax}({\mathbf{QK}}_{i,j}^T)\text{，在 }i,j\text{ 上迭代}$$

  其中 $s_{ij}^a$和 $s_{ij}^b$ 是随着每个 $i$ 和 $j$迭代更新的 softmax 统计归一化值。最终，通过跟踪 softmax 统计归一化值再加上一些聪明的数学技巧，与默认的自注意力层相比，Flash 注意力的计算结果 完全相同，而内存成本仅随着 $N$ 线性增加。如果想要深入理解，可以阅读 Flash Attention 的论文。

  仅看这个公式，Flash 注意力需要不断重新计算softmax 统计归一化值，需要更多的 FLOP，直觉上应该比默认的自注意力公式要慢很多。但实际上，因为使用了更快的片上内存 (SRAM)，Flash 注意力的推理速度要快得多。

6.1.2 GPU 硬件分析

  上一节讲到因为使用了更快的片上内存 (SRAM)，Flash 注意力比标准的注意力计算快很多，所以了解GPU内存和各种操作的性能特征是很有必要的。

  以 A100 (40GB HBM) 为例，下面显示其内存层次结构的粗略图。SRAM内存分布在108个流式多处理器(SMs)上，每个处理器192KB。从图中可以看到，片上SRAM读写操作比HBM快得多，但比HBM小得多。GPU 的典型操作方式是使用大量的线程来执行一个操作，这个操作被称为内核。输入从HBM加载到寄存器和SRAM，并在计算后写回HBM。在计算方面，使用Tensor Core的BFLOAT16 的理论峰值吞吐量为 312 TFLOPS。

  算法对于内存带宽的需求通常使用 计算强度 (arithmetic intensity) 来表示，单位是 OPs/byte，意思是在算法中平均每读入单位数据，能支持多少次运算操作。它有助于理解操作的瓶颈，即计算约束(Compute-bound)或带宽约束(Bandwidth-bound, or Memory-bound)。

• 算力 $\pi$ ：也称为计算平台的性能上限，指的是一个计算平台倾尽全力每秒钟所能完成的浮点运算数。单位是 FLOPS or FLOP/s。
• 带宽 $\beta$ ：也即计算平台的带宽上限，指的是一个计算平台倾尽全力每秒所能完成的内存交换量。单位是Byte/s。
• 计算强度上限$I_{max}=\frac{\pi }{\beta }$ ：两个指标相除即可得到计算平台的计算强度上限。它描述的是在这个计算平台上，单位内存交换最多用来进行多少次计算。单位是FLOPs/Byte。
• 模型的理论性能 $P$： 我们最关心的指标，即模型在计算平台上所能达到的每秒浮点运算次数（理论值）。单位是FLOPSorFLOP/s。

如下图所示，Roof-line 描述了模型在一个计算平台的限制下，到底能达到多快的浮点计算速度，即算力决定“屋顶”的高度（绿色线段），带宽决定“房檐”的斜率（红色线段）。

Roof-line 划分出的两个瓶颈区域，即

$$P=\{\begin{array}{ll}\beta \cdot I,& whenI<I_{max}\text{Memory Bound}\\ \pi ,& whenI⩾I_{max}\text{Compute Bound}\end{array}$$

• 计算约束——此时HBM访问所花费的时间相对较低，不管模型的计算强度$I$ 有多大，它的理论性能 $P$ 最大只能等于计算平台的算力$\pi$。例如，具有较大内维数的矩阵乘法和具有大量通道的卷积。
• 带宽约束——当模型的计算强度 $I$ 小于计算平台的计算强度上限$I_{max}$ 时，由于此时模型位于“房檐”区间，因此模型理论性能 $P$ 的大小完全由计算平台的带宽上限$\beta$（房檐的斜率）以及模型自身的计算强度$I$ 所决定。例如，elementwise 操作 (如activation, dropout 等) 和 规约操作 (如sum, softmax, batch normalization, layer normalization等)。

  在 self-attention 中，计算速度比内存速度快得多，进程(操作)越来越多地受到内存(HBM)访问的瓶颈。因此，FlashAttention论文的目标是尽可能高效地使用SRAM来加快计算速度。

6.1.3 算法

  从之前对标准attention计算公式的分析可以知道，复杂度为$O(N^2)$的矩阵对HBM及其重复读写是一个主要瓶颈，要解决这个问题，需要做两件主要的事情：

• 在不访问整个输入的情况下计算 softmax
• 不为反向传播存储大的中间 attention 矩阵

为此 FlashAttention 提出了两种方法来分布解决上述问题：tiling 和 recomputation。

• tiling - 注意力计算被重新构造，将输入分割成块，并通过在输入块上进行多次传递来递增地执行softmax操作。
• recomputation - 存储来自前向的 softmax 归一化因子，以便在反向中快速重新计算芯片上的 attention，这比从HBM读取中间矩阵的标准注意力方法更快

  关于 tiling和 recomputation部分的公式推导，可参考原论文，或博客《从 FlashAttention 到 PagedAttention, 如何进一步优化 Attention 性能》

  由于重新计算，这确实导致FLOPs增加，但是由于大量减少HBM访问，FlashAttention运行速度更快(在GPT-2上高达7.6倍)。

  总的来说，Flash 注意力的主要思想是分割输入，将它们从慢速HBM加载到快速SRAM，然后计算这些块的 attention 输出。在将每个块的输出相加之前，将其按正确的归一化因子进行缩放，从而得到正确的结果。因为所有中间写入和读取操作都可以使用快速片上SRAM 来完成，而不必访问慢速显存HBM，所以该算法在数学上给出相同输出的同时，速度更快且内存效率更高。所以如果能用的话，我们没有理由不用 Flash 注意力。

6.2 FlashAttention-2

6.2.1 FlashAttention-2简介

FlashAttention-2 是标准注意力机制的更快、更高效的实现，可以通过以下方式显著加快推理速度：

1. 在序列长度上并行化注意力计算
2. 在GPU线程之间分配工作，以减少它们之间的通信和共享内存读/写

FlashAttention-2现在支持Llama、Mistral和Falcon模型的推理，在开始之前，请确保已安装FlashAttention-2

pip install flash-attn --no-build-isolation

要启用FlashAttention-2，请在from_pretrained()中添加use_flash_attention_2参数。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM

model_id = "tiiuae/falcon-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16, 
    use_flash_attention_2=True,)

  FlashAttention-2 只能在模型的 dtype 为 fp16 或 bf16 时使用，并且它只能在 Nvidia GPU 上运行。在使用 FlashAttention-2 之前，请确保将模型转换为适当的 dtype，并将其加载到受支持的设备上。

  FlashAttention-2 还可以与其他优化技术（如量化）结合使用，以进一步加快推理速度。例如，您可以将 FlashAttention-2 与 8 位或 4 位量化结合使用：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM

model_id = "tiiuae/falcon-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# load in 8bit
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    load_in_8bit=True,
    use_flash_attention_2=True,)

# load in 4bit
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    load_in_4bit=True,
    use_flash_attention_2=True,)

6.2.2 加速效果

  使用FlashAttention-2在推理过程中可以获得显著的速度提升，特别是对于具有长序列输入的情况。然而，FlashAttention-2不支持在包含填充标记（padding tokens）的序列上计算注意力分数，因此在批处理推理中，如果序列包含填充标记，您需要手动进行填充和取消填充操作，这将导致批量生成包含填充标记的情况下速度显著降低。

  为了解决这个问题，建议在训练过程中使用FlashAttention-2，但在序列中不使用填充标记，可以通过对数据集进行打包或连接序列（concatenating sequences）直到达到最大序列长度来实现。

1. 对于包含填充标记的序列（生成时包含填充标记），您需要取消填充，以正确计算注意力分数。在相对较小的序列长度下，单次前向传递会产生一些开销，从而导致小幅度的加速（在下面的示例中，输入的30%被填充标记占据）：
   

2. 但对于更大的序列长度，FlashAttention加速效果更明显
   

3. 对于tiiuae/falcon-7b的单次前向传递，当序列长度为4096且没有填充标记时，预期的加速情况如下：
   

4. 对于meta-llama/Llama-7b-hf的单次前向传递，当序列长度为4096且不包含填充标记时，预期的加速情况如下：
   

5. 在A100和H100上，Head dimension为64或128，hidden dimension=2048（也就是32或16个head），batch size=16k / seqlen时，加速效果如下：
   
   

6. 下图展示了展示了使用FlashAttention时的内存节省情况。
   因为标准注意力的内存消耗与序列长度的平方成正比，而FlashAttention的内存消耗与序列长度成线性关系，所以序列越长，使用FlashAttention越省内存。在序列长度为2K时，内存节省为10倍，而在4K时为20倍，这意味着使用FlashAttention后，你可以在更大的序列长度上进行训练。
   

更多详细信息请查看flash-attention。

七、组合优化（GPU）

  上述几种优化技术可以组合使用，以获得最佳推理性能。例如，您可以加载 4 位模型，然后启用FlashAttention和BetterTransformer：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# load model in 4-bit
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", quantization_config=quantization_config)

# enable BetterTransformer
model = model.to_bettertransformer()

input_text = "Hello my dog is cute and"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda")

# enable FlashAttention
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    outputs = model.generate(**inputs)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

其它优化方法，可参考知乎博客《LLM 的推理优化技术纵览》。