大模型推理

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)

後續步驟

有關大模型推理的更詳細解釋，請務必檢視概念指南！