大模型推理

Accelerate 提供的最大改进之一是大模型推理，它允许您对无法完全适应显卡的模型执行推理。

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

加速

加载 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 等其他工具来正常启动您的脚本。

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

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)

后续步骤

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