Transformers 文档

GPU 推理

Transformers

加入 Hugging Face 社区

并获得增强版文档体验

协作模型、数据集和 Spaces

使用加速推理更快地运行示例

在文档主题之间切换

开始使用

GPU 推理

与 CPU 不同，GPU 是机器学习的标准硬件选择，因为它们针对内存带宽和并行性进行了优化。为了跟上现代模型的更大规模，或者在现有和较旧的硬件上运行这些大型模型，您可以使用多种优化方法来加速 GPU 推理。在本指南中，您将学习如何使用 FlashAttention-2（一种更节省内存的注意力机制）、BetterTransformer（一种 PyTorch 原生快速路径执行）以及 bitsandbytes 对您的模型进行低精度量化。最后，了解如何使用 🤗 Optimum 使用 ONNX Runtime 在 Nvidia 和 AMD GPU 上加速推理。

此处描述的大多数优化也适用于多 GPU 设置！

FlashAttention-2

FlashAttention-2 处于实验阶段，在未来版本中可能会发生很大变化。

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

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

FlashAttention-2 目前支持以下架构：

您可以通过在 GitHub 上提交 Issue 或 Pull Request 来请求为其他模型添加 FlashAttention-2 支持。

在开始之前，请确保您已安装 FlashAttention-2。

NVIDIA

AMD

要启用 FlashAttention-2，请将参数 attn_implementation="flash_attention_2" 传递给 from_pretrained()

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,
    attn_implementation="flash_attention_2",
)

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

您也可以设置 use_flash_attention_2=True 来启用 FlashAttention-2，但它已被弃用，建议使用 attn_implementation="flash_attention_2"。

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,
    attn_implementation="flash_attention_2",
)

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

预期加速

您可以从推理中获得显著的加速，特别是对于具有长序列的输入。但是，由于 FlashAttention-2 不支持使用填充标记计算注意力分数，因此当序列包含填充标记时，您必须手动填充/取消填充用于批量推理的注意力分数。这会导致使用填充标记进行批量生成的显著减速。

为了克服这一点，您应该在训练期间（通过打包数据集或连接序列，直到达到最大序列长度）在序列中不使用填充标记的情况下使用 FlashAttention-2。

对于 tiiuae/falcon-7b 上的单个前向传递，序列长度为 4096，各种批次大小不包含填充标记，预期的加速为

对于 meta-llama/Llama-7b-hf 上的单个前向传递，序列长度为 4096，各种批次大小不包含填充标记，预期的加速为

对于包含填充标记的序列（使用填充标记生成），您需要取消填充/填充输入序列以正确计算注意力分数。对于相对较小的序列长度，单个前向传递会产生开销，导致加速较小（在以下示例中，30% 的输入填充了填充标记）。

但对于更大的序列长度，您可以预期获得更大的加速收益。

FlashAttention 更加内存效率，这意味着您可以在更大的序列长度上进行训练，而不会遇到内存不足的问题。对于更大的序列长度，您可能可以将内存使用量减少高达 20 倍。请查看 flash-attention 存储库以获取更多详细信息。

PyTorch 缩放点积注意力

PyTorch 的 torch.nn.functional.scaled_dot_product_attention (SDPA) 也可以在后台调用 FlashAttention 和内存高效的注意力内核。SDPA 支持目前正在 Transformers 中原生添加，并且在 torch>=2.1.1 中默认使用，前提是实现可用。您也可以在 from_pretrained() 中设置 attn_implementation="sdpa" 以显式请求使用 SDPA。

目前，Transformers 支持以下架构的 SDPA 推理和训练。

FlashAttention 只能用于具有 fp16 或 bf16 火炬类型的模型，因此请确保首先将您的模型转换为适当的类型。内存高效的注意力后端能够处理 fp32 模型。

SDPA 不支持某些注意力参数集，例如 head_mask 和 output_attentions=True。在这种情况下，您应该看到警告消息，我们将回退到（较慢的）急切执行。

默认情况下，SDPA 会选择性能最高的可用内核，但您可以使用 torch.backends.cuda.sdp_kernel 作为上下文管理器来检查给定设置（硬件、问题大小）中是否可用后端。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", torch_dtype=torch.float16).to("cuda")

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

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

如果您看到以下回溯错误，请尝试使用 PyTorch 的夜间版本，该版本可能对 FlashAttention 的覆盖范围更广。

RuntimeError: No available kernel. Aborting execution.

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

BetterTransformer

BetterTransformer 的一些功能正在上游到 Transformers，并默认支持原生 torch.nn.scaled_dot_product_attention。BetterTransformer 的覆盖范围仍然比 Transformers SDPA 集成更广，但您可以预期 Transformers 中越来越多的架构将原生支持 SDPA。

查看我们在使用 PyTorch 2.0 对 🤗 解码器模型进行开箱即用的加速和内存节省中使用 BetterTransformer 和缩放点积注意力的基准测试，并了解有关 BetterTransformer 博客文章中的快速路径执行的更多信息。

BetterTransformer 通过其快速路径（Transformer 函数的原生 PyTorch 专用实现）执行来加速推理。快速路径执行中的两个优化是

融合，它将多个顺序操作组合成一个“内核”，以减少计算步骤的数量
跳过填充标记的固有稀疏性，以避免使用嵌套张量进行不必要的计算

BetterTransformer 还将所有注意力操作转换为使用内存效率更高的缩放点积注意力 (SDPA)，并且它在后台调用了优化的内核，例如 FlashAttention。

在开始之前，请确保您已安装 🤗 Optimum installed。

然后，您可以使用 PreTrainedModel.to_bettertransformer() 方法启用 BetterTransformer。

model = model.to_bettertransformer()

您可以使用 reverse_bettertransformer() 方法返回原始 Transformers 模型。在将模型保存为使用规范 Transformers 模型之前，您应该使用此方法。

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

bitsandbytes

bitsandbytes 是一个量化库，它包括对 4 位和 8 位量化的支持。量化会减少模型相对于其原生全精度版本的大小，使其更容易将大型模型适应内存有限的 GPU。

请确保您已安装 bitsandbytes 和 🤗 Accelerate。

# these versions support 8-bit and 4-bit
pip install bitsandbytes>=0.39.0 accelerate>=0.20.0

# install Transformers
pip install transformers

4 位

要在推理时以 4 位加载模型，请使用 load_in_4bit 参数。device_map 参数是可选的，但我们建议将其设置为 "auto"，以便 🤗 Accelerate 可以根据环境中可用的资源自动并有效地分配模型。

from transformers import AutoModelForCausalLM

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

要在推理时以 4 位加载模型并使用多个 GPU，您可以控制要分配给每个 GPU 的 GPU RAM 量。例如，要将 600MB 内存分配给第一个 GPU，并将 1GB 内存分配给第二个 GPU

max_memory_mapping = {0: "600MB", 1: "1GB"}
model_name = "bigscience/bloom-3b"
model_4bit = AutoModelForCausalLM.from_pretrained(
    model_name, device_map="auto", load_in_4bit=True, max_memory=max_memory_mapping
)

8 位

如果您好奇并有兴趣了解更多关于 8 位量化背后的概念，请阅读使用 Hugging Face Transformers、Accelerate 和 bitsandbytes 大规模进行 8 位矩阵乘法变换的简明介绍博客文章。

要在推理时以 8 位加载模型，请使用 load_in_8bit 参数。device_map 参数是可选的，但我们建议将其设置为 "auto"，以便 🤗 Accelerate 可以根据环境中可用的资源自动并有效地分配模型。

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

model_name = "bigscience/bloom-2b5"
model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))

如果您要以 8 位加载模型进行文本生成，则应使用 generate() 方法，而不是 Pipeline 函数，后者没有针对 8 位模型进行优化，速度会更慢。一些采样策略，例如核采样，也不受 Pipeline 的支持，用于 8 位模型。您还应将所有输入放在与模型相同的设备上。

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_name = "bigscience/bloom-2b5"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model_8bit = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=BitsAndBytesConfig(load_in_8bit=True))

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)

要在推理时以 4 位加载模型并使用多个 GPU，您可以控制要分配给每个 GPU 的 GPU RAM 量。例如，要将 1GB 内存分配给第一个 GPU，并将 2GB 内存分配给第二个 GPU

max_memory_mapping = {0: "1GB", 1: "2GB"}
model_name = "bigscience/bloom-3b"
model_8bit = AutoModelForCausalLM.from_pretrained(
    model_name, device_map="auto", load_in_8bit=True, max_memory=max_memory_mapping
)

您可以尝试运行 110 亿参数的 T5 模型或 30 亿参数的 BLOOM 模型，在 Google Colab 的免费层 GPU 上进行推理！

🤗 Optimum

在 NVIDIA GPU 上加速推理和 AMD GPU 上加速推理指南中了解有关在 🤗 Optimum 中使用 ORT 的更多详细信息。本节仅提供一个简短的示例。

ONNX 运行时 (ORT) 是一种模型加速器，它支持在使用 ROCm 堆栈的 Nvidia GPU 和 AMD GPU 上加速推理。ORT 使用诸如将常见操作融合到单个节点和常量折叠之类的优化技术，以减少执行的计算次数并加快推理速度。ORT 还将计算量最大的操作放在 GPU 上，并将其余操作放在 CPU 上，以便在两个设备之间智能地分配工作负载。

ORT 受 🤗 Optimum 支持，它可以在 🤗 Transformers 中使用。您需要使用 ORTModel 来完成您正在解决的任务，并指定 provider 参数，该参数可以设置为 CUDAExecutionProvider、ROCMExecutionProvider 或 TensorrtExecutionProvider。如果您想加载尚未导出到 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",
)

现在您可以随意使用模型进行推理。

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

结合优化

通常可以结合上面描述的几种优化技术，以获得模型可能的最佳推理性能。例如，您可以以 4 位加载模型，然后启用 BetterTransformer 与 FlashAttention

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

< > 在 GitHub 上更新

←CPU 推理实例化大型模型→