Accelerate 文档

大模型推理

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

大模型推理

Accelerate 提供的最大进步之一是大模型推理,它允许你对模型进行推理,即使模型无法完全放入你的显卡。

本教程将向你展示如何在 Accelerate 和 Hugging Face 生态系统中使用大模型推理。

Accelerate

下面展示了加载 PyTorch 模型的典型工作流程。ModelClass 是一个超出你的设备 (mps 或 cuda 或 xpu) GPU 内存的模型。

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

使用大模型推理,第一步是使用 init_empty_weights 上下文管理器初始化模型的空骨架。这不需要任何内存,因为 my_model 是“无参数的”。

from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)

接下来,权重被加载到模型中进行推理。

load_checkpoint_and_dispatch() 方法将检查点加载到你的空模型中,并将每个层的权重分发到所有可用设备上,首先从最快的设备 (GPU、MPS、XPU、NPU、MLU、SDAA、MUSA) 开始,然后再移动到较慢的设备 (CPU 和硬盘驱动器)。

设置 device_map="auto" 会自动首先填充 GPU 上的所有可用空间,然后是 CPU,最后是硬盘驱动器 (绝对最慢的选项),如果仍然没有足够的内存。

有关如何设计你自己的设备映射的更多详细信息,请参阅 设计设备映射 指南。

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

如果有某些不应拆分的层“块”,请将它们传递给 `no_split_module_classes` (有关更多详细信息,请参阅 此处)。

模型的权重也可以被分片到多个检查点以节省内存,例如当 `state_dict` 不适合内存时 (有关更多详细信息,请参阅 此处)。

现在模型已完全分发,你可以执行推理。

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

每次输入通过一个层时,它都会从 CPU 发送到 GPU (或从磁盘到 CPU 到 GPU),计算输出,然后将该层从 GPU 中移除,向下传递。虽然这会给推理增加一些开销,但它使你能够在你的系统上运行任何大小的模型,只要最大的层适合你的 GPU。

可以使用多个 GPU 或“模型并行”,但任何给定时刻只有一个 GPU 处于活动状态。这迫使 GPU 等待前一个 GPU 将输出发送给它。你应该使用 Python 正常启动你的脚本,而不是像 torchrun 和 accelerate launch 这样的其他工具。

你可能还会对*管线并行*感兴趣,它一次性利用所有可用的 GPU,而不是一次只激活一个 GPU。尽管这种方法不太灵活。有关更多详细信息,请参阅 内存高效的管线并行 指南。

请看下面大模型推理的完整示例。

import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

with init_empty_weights():
    model = MyModel(...)

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Hugging Face 生态系统

Hugging Face 生态系统中的其他库,如 Transformers 或 Diffusers,在其 `from_pretrained` 构造函数中支持大模型推理。

你只需要在 `from_pretrained` 中添加 `device_map="auto"` 即可启用大模型推理。

例如,使用大模型推理加载 Big Sciences T0pp 110 亿参数模型。

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")

加载模型后,之前提到的空初始化和智能分发步骤会被执行,模型已完全准备好使用你机器中的所有资源。通过这些构造函数,你还可以通过指定 `torch_dtype` 参数以较低的精度加载模型,从而节省更多内存。

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto", torch_dtype=torch.float16)

下一步

有关大模型推理的更详细解释,请务必查看 概念指南

< > 在 GitHub 上更新