Skip to content

大模型推理

Accelerate 提供的最大进步之一是 大模型推理,它允许你在模型无法完全适应你的显卡时进行推理。

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

Accelerate

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

py
import torch

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

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

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

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

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

设置 device_map="auto" 会自动先填满所有可用的 GPU 内存,然后是 CPU,最后如果仍然内存不足,则会使用硬盘(最慢的选项)。

TIP

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

py
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 无法适应内存时(详情请参见这里)。

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

py
input = torch.randn(2,3)
input = input.to("cuda")
output = model(input)

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

多个 GPU,或称为“模型并行”,可以被利用,但任何时候只有一个 GPU 是活跃的。这迫使当前的 GPU 必须等待前一个 GPU 发送输出。你应该使用 Python 正常启动脚本,而不是使用其他工具如 torchrun 和 accelerate launch。

TIP

你可能也对 管道并行 感兴趣,它利用所有可用的 GPU,而不仅仅是一次只有一个 GPU 活跃。这种方法的灵活性较低。更多详情,请参阅 内存高效的管道并行 指南。

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

py
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)
input = input.to("cuda")
output = model(input)

Hugging Face 生态系统

Hugging Face 生态系统中的其他库,如 Transformers 或 Diffusers,在它们的 [~transformers.PreTrainedModel.from_pretrained] 构造函数中支持大模型推理。

你只需要在 [~transformers.PreTrainedModel.from_pretrained] 中添加 device_map="auto" 来启用大模型推理。

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

py
from transformers import AutoModelForSeq2SeqLM

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

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

py
from transformers import AutoModelForSeq2SeqLM

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

下一步

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