如何使用 Optimum 和 BetterTransformer？

安装依赖项

您可以轻松地使用 BetterTransformer 与 🤗 Optimum 集成，首先安装以下依赖项

pip install transformers accelerate optimum

此外，请确保按照 PyTorch 官方网站上的指南安装最新版本的 PyTorch。请注意，BetterTransformer API 仅与 torch>=1.13 兼容，因此在开始之前请确保在您的环境中安装了此版本。如果您想从 scaled_dot_product_attention 函数（用于基于解码器的模型）中获益，请确保使用至少 torch>=2.0。

步骤 1：加载您的模型

首先，使用 🤗 Transformers 加载您的 Hugging Face 模型。请确保下载 BetterTransformer API 支持的模型之一

>>> from transformers import AutoModel

>>> model_id = "roberta-base"
>>> model = AutoModel.from_pretrained(model_id)

有时，您可以使用 accelerate 库直接在您的 GPU 设备上加载您的模型，因此您可以选择尝试以下命令

>>> from transformers import AutoModel

>>> model_id = "roberta-base"
>>> model = AutoModel.from_pretrained(model_id, device_map="auto")

步骤 2：将您的模型设置为首选设备

如果您没有使用 device_map="auto" 来加载您的模型（或者您的模型不支持 device_map="auto"），您可以手动将您的模型设置为 GPU

>>> model = model.to(0) # or model.to("cuda:0")

步骤 3：将您的模型转换为 BetterTransformer！

现在是时候使用 BetterTransformer API 转换您的模型了！您可以运行以下命令

>>> from optimum.bettertransformer import BetterTransformer

>>> model = BetterTransformer.transform(model)

默认情况下，BetterTransformer.transform 会覆盖您的模型，这意味着您之前使用的原生模型将无法再使用。如果您想出于某种原因保留它，只需添加标志 keep_original_model=True！

>>> from optimum.bettertransformer import BetterTransformer

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

如果您的模型不支持 BetterTransformer API，这将在错误跟踪中显示。还要注意，基于解码器的模型（OPT、BLOOM 等）目前尚不支持，但这是 PyTorch 未来路线图的一部分。

管道兼容性

Transformer 的管道也与该集成兼容，您可以使用 BetterTransformer 作为管道的加速器。以下代码段显示了如何操作

>>> from optimum.pipelines import pipeline

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

如果您想在 GPU 设备上运行管道，请运行

>>> from optimum.pipelines import pipeline

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

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

>>> from transformers import pipeline

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

有关进一步使用，请参阅 pipeline 的官方文档。如果您遇到任何问题，请随时在 GitHub 上打开问题！

训练兼容性

您现在可以从 BetterTransformer API 中获益，用于您的训练脚本。只需确保在保存模型之前通过调用 BetterTransformer.reverse 将模型转换回其原始版本。以下代码段显示了如何操作

from optimum.bettertransformer import BetterTransformer
from transformers import AutoModelForCausalLM

with torch.device(“cuda”):
    model = AutoModelForCausalLM.from_pretrained(“gpt2-large”, torch_dtype=torch.float16)

model = BetterTransformer.transform(model)

# do your inference or training here

# if training and want to save the model
model = BetterTransformer.reverse(model)
model.save_pretrained("fine_tuned_model")
model.push_to_hub("fine_tuned_model")

< > 在 GitHub 上更新