如何使用 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 未来的路线图。

Pipeline 兼容性

Transformer 的 pipeline 也与此集成兼容，您可以将 BetterTransformer 用作 pipeline 的加速器。以下代码片段展示了如何操作

>>> from optimum.pipelines import pipeline

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

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

>>> 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 上提出 issue！

训练兼容性

您现在可以在您的训练脚本中受益于 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 上更新