GPU

GPU 是机器学习的标准硬件，因为它们针对内存带宽和并行性进行了优化。随着现代模型规模的不断扩大，确保 GPU 能够高效处理并提供最佳性能变得前所未有的重要。

本指南将演示几种优化 GPU 推理的方法。下面所示的优化方法可以相互结合，以实现更好的性能，并且它们也适用于分布式 GPU。

bitsandbytes

bitsandbytes 是一个支持 8 位和 4 位量化的量化库。量化表示与原始全精度格式相比，权重具有较低的精度。它减少了内存需求，并使大型模型更容易适应内存。

首先确保安装了 bitsandbytes 和 Accelerate。

pip install bitsandbytes accelerate

8位

4位

Optimum

Optimum 是一个 Hugging Face 库，专注于优化各种硬件上的模型性能。它支持 ONNX Runtime (ORT)，这是一种模型加速器，适用于包括 NVIDIA GPU 和使用 ROCm 堆栈的 AMD GPU 在内的各种硬件和框架。

ORT 使用优化技术，将常见操作融合到单个节点中，并进行常量折叠以减少计算次数。ORT 还将计算量最大的操作放在 GPU 上，其余操作放在 CPU 上，以智能地分配两个设备之间的工作负载。

Optimum 提供了 ORTModel 类用于加载 ONNX 模型。根据下表设置 provider 参数。

提供者	硬件
CUDAExecutionProvider	支持 CUDA 的 GPU
ROCMExecutionProvider	AMD Instinct、Radeon Pro、Radeon GPU
TensorrtExecutionProvider	TensorRT

例如，加载 distilbert/distilbert-base-uncased-finetuned-sst-2-english 检查点进行序列分类。此检查点包含一个 model.onnx 文件。如果检查点没有 model.onnx 文件，请设置 export=True 以动态将检查点转换为 ONNX 格式。

from optimum.onnxruntime import ORTModelForSequenceClassification

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

现在您可以在 Pipeline 中使用模型进行推理。

from optimum.pipelines import pipeline
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert/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.")

在 NVIDIA GPU 上的加速推理和 AMD GPU 上的加速推理指南中了解更多关于将 ORT 与 Optimum 结合使用的详细信息。

BetterTransformer

BetterTransformer 是直接在 GPU 等硬件级别上执行专门的 Transformers 函数的 *快速路径*。快速路径执行主要包含两个部分。

将多个操作融合到一个内核中，以实现更快、更高效的执行
通过嵌套张量跳过填充令牌的不必要计算

一些 BetterTransformer 功能正在被引入 Transformers，默认支持原生 torch.nn.functional.scaled_dot_product_attention (SDPA)。BetterTransformer 比 Transformers SDPA 集成具有更广泛的覆盖范围，但您可以预期越来越多的架构将原生支持 Transformers 中的 SDPA。

BetterTransformer 通过 Optimum 提供，使用 to_bettertransformer()。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom")
model = model.to_bettertransformer()

调用 reverse_bettertransformer() 并先保存它，以将模型恢复到原始的 Transformers 模型。

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

请参阅 PyTorch 2.0 中 🤗 解码器模型的开箱即用加速和内存节省中的 BetterTransformer 和缩放点积注意力性能基准。如果您有兴趣了解更多，BetterTransformer 博客文章也更详细地讨论了快速路径执行。

缩放点积注意力 (SDPA)

PyTorch 的 torch.nn.functional.scaled_dot_product_attention (SDPA) 是缩放点积注意力机制的原生实现。SDPA 是 transformer 模型中使用的注意力机制的一种更高效、更优化的版本。

有三种受支持的实现可用。

FlashAttention2 仅支持 fp16 或 bf16 torch 类型的模型。请务必先将模型转换为适当的类型。
xFormers 或内存高效注意力能够支持 fp32 torch 类型的模型。
缩放点积注意力的 C++ 实现

当有可用的实现时，SDPA 默认用于 PyTorch v2.1.1 及更高版本。您可以通过在 from_pretrained() 中设置 attn_implementation="sdpa" 来显式启用 SDPA。某些注意力参数，例如 head_mask 和 output_attentions=True，不受支持，并会返回警告，指出 Transformers 将回退到（较慢的）eager 实现。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", attn_implementation="sdpa")

SDPA 选择可用的性能最佳的实现，但您也可以使用 torch.nn.attention.sdpa_kernel 作为上下文管理器显式选择实现。下面的示例展示了如何使用 enable_flash=True 启用 FlashAttention2 实现。

import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto").to("cuda")

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

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    outputs = model.generate(**inputs)

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

如果您遇到以下 RuntimeError，请尝试安装 PyTorch 的每夜版本，它对 FlashAttention 具有更广泛的覆盖范围。

RuntimeError: No available kernel. Aborting execution.

pip3 install -U --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118

FlashAttention

FlashAttention 也可作为独立包使用。它可以通过以下方式显著加速推理：

额外地对序列长度上的注意力计算进行并行化
在 GPU 线程之间划分工作，以减少它们之间的通信和共享内存读写

首先为您正在使用的硬件安装 FlashAttention。

NVIDIA

AMD

通过在 from_pretrained() 中设置 attn_implementation="flash_attention_2" 来启用 FlashAttention2。FlashAttention2 仅支持 fp16 或 bf16 torch 类型的模型。请务必先将模型转换为适当的数据类型。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")

基准测试

FlashAttention2 显著加速了推理，特别是对于长序列的输入。然而，由于 FlashAttention2 不支持使用填充令牌计算注意力分数，如果序列包含填充令牌，则必须手动对批量推理的注意力分数进行填充和取消填充。缺点是使用填充令牌进行批量生成速度较慢。

短序列长度

长序列长度

为了避免这种减速，在训练期间使用 FlashAttention2，序列中不带填充令牌。打包数据集或连接序列直到达到最大序列长度。

tiiuae/falcon-7b

meta-llama/Llama-7b-hf

< > 在 GitHub 上更新