Hugging Face 的徽标

Transformers 文档

GPU

Hugging Face's logo
加入 Hugging Face 社区

并获取增强的文档体验

开始

GPU

GPU 是机器学习的标准硬件，因为它们针对内存带宽和并行性进行了优化。随着现代模型尺寸的不断增大，确保 GPU 能够高效处理并提供最佳性能比以往任何时候都更加重要。

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

bitsandbytes

bitsandbytes 是一个量化库，支持 8 位和 4 位量化。量化以低于原始全精度格式的精度表示权重。它可以减少内存需求，并使大型模型更容易放入内存。

请确保先安装 bitsandbytes 和 Accelerate。

pip install bitsandbytes accelerate
8 位
4 位

对于使用 8 位量化进行文本生成，您应该使用 generate() 而不是高级 Pipeline API。Pipeline 返回的性能较慢，因为它没有针对 8 位模型进行优化，并且某些采样策略（核采样）也不受支持。

设置一个 BitsAndBytesConfig 并设置 load_in_8bit=True 以加载 8 位精度的模型。BitsAndBytesConfig 将传递给 from_pretrained() 中的 quantization_config 参数。

通过设置 device_map=“auto”，允许 Accelerate 自动将模型分布在您可用的硬件上。

将所有输入放在与模型相同的设备上。

from transformers import BitsAndBytesConfig, AutoTokenizer, AutoModelForCausalLM

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
tokenizer = AutoTokenizer("meta-llama/Llama-3.1-8B")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", quantization_config=quantization_config)

prompt = "Hello, my llama is cute"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
generated_ids = model.generate(**inputs)
outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

对于分布式设置，使用 max_memory 参数创建每个 GPU 要分配的内存量映射。下面的示例将 16GB 内存分配给第一个 GPU，将 16GB 内存分配给第二个 GPU。

max_memory_mapping = {0: "16GB", 1: "16GB"}
model_8bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B", device_map="auto", quantization_config=quantization_config, max_memory=max_memory_mapping
)

在《使用 Hugging Face Transformers、Accelerate 和 bitsandbytes 大规模变压器的 8 位矩阵乘法入门指南》博客文章中，详细了解 8 位量化的基本概念。

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")

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

缩放点积注意力 (SDPA)

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

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

  • FlashAttention2 仅支持 fp16 或 bf16 torch 类型的模型。请确保先将模型转换为适当的类型。
  • xFormers 或 Memory-Efficient Attention 能够支持 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 的 nightly 版本，该版本对 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 也可作为独立软件包使用。它可以通过以下方式显着加速推理：

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

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

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

通过在 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 不支持使用填充标记计算注意力分数，因此如果序列包含填充标记，则必须手动填充和取消填充批量推理的注意力分数。缺点是使用填充标记的批量生成速度较慢。

短序列长度
长序列长度

对于相对较小的序列长度，单个前向传递会产生开销，从而导致较小的加速。下图显示了使用填充的 meta-llama/Llama-7b-hf 进行单次前向传递的预期加速。

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

tiiuae/falcon-7b
meta-llama/Llama-7b-hf

下图显示了使用 tiiuae/falcon-7b、序列长度为 4096 且批量大小不同的情况下，不带填充标记的单次前向传递的预期加速。

< > 更新 on GitHub

←torch.compile 分布式 GPU 推理→