加载模型

Transformers 提供了许多预训练模型，只需一行代码即可使用。它需要一个模型类以及 from_pretrained() 方法。

调用 from_pretrained() 以下载并加载存储在 Hugging Face Hub 上的模型权重和配置。

The from_pretrained() method loads weights stored in the safetensors file format if they’re available. Traditionally, PyTorch model weights are serialized with the pickle utility which is known to be insecure. Safetensor files are more secure and faster to load.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")

本指南介绍了模型的加载方式、加载模型的多种方法、如何克服超大模型的内存问题，以及如何加载自定义模型。

模型与配置

所有模型都有一个 configuration.py 文件，其中包含特定属性，例如隐藏层数、词汇表大小、激活函数等。您还会找到一个 modeling.py 文件，它定义了各层以及每层内部执行的数学运算。modeling.py 文件会读取 configuration.py 中的模型属性，并据此构建模型。此时，您获得的是一个具有随机权重的模型，需要经过训练才能输出有意义的结果。

架构指的是模型的骨架，而 检查点指的是给定架构下模型的权重。例如，BERT 是一种架构，而 google-bert/bert-base-uncased 是一个检查点。你会看到模型一词与架构和检查点互换使用。

您可以加载两种通用类型的模型：

一个仅输出隐藏状态的裸模型，例如 AutoModel 或 LlamaModel。
附加了特定 head 的模型，例如 AutoModelForCausalLM 或 LlamaForCausalLM，用于执行特定任务。

模型类

要获取预训练模型，需要将权重加载到模型中。这通过调用 from_pretrained() 完成，该方法接受来自 Hugging Face Hub 或本地目录的权重。

有两个模型类，分别是 AutoModel 类和特定于模型的类。

AutoModel

特定于模型的类

大型模型

大型预训练模型在加载时需要占用大量内存。加载过程包括：

创建具有随机权重的模型
加载预训练权重
将预训练权重加载到模型上

您需要足够的内存来容纳两份模型权重（随机初始化和预训练），具体取决于您的硬件配置，这可能无法实现。在分布式训练环境中，这一挑战更为严峻，因为每个进程都会加载一个预训练模型。

Transformers 通过快速初始化、分片检查点、Accelerate 的 Big Model Inference 功能以及对低比特数据类型的支持，缓解了其中一些与内存相关的挑战。

分片检查点

save_pretrained() 会自动对大于 50GB 的检查点进行分片。这有助于减少大模型的分片数量，并简化文件管理。

参数并行加载，峰值内存仅取决于模型大小。在 save_pretrained() 中使用 max_shard_size 来设置分片前的最大检查点大小。

需要动态权重转换的模型的内存占用取决于模型大小以及单次转换中最大参数的大小。这通常适用于混合专家（MoE）模型，其内存占用为模型大小加上单层中的专家数量。请参阅动态权重加载器指南，以了解模型加载的更多详细信息。

save_pretrained() 还会创建一个索引文件，用于将参数名称映射到其分片文件。该索引包含两个键：metadata 和 weight_map。

import json
import os
import tempfile

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="50GB")
    with open(os.path.join(tmp_dir, "model.safetensors.index.json"), "r") as f:
        index = json.load(f)

print(index.keys())

metadata 存储模型的总大小。

index["metadata"]
{'total_size': 28966928384}

weight_map 将每个参数映射到其分片文件。

index["weight_map"]
{'lm_head.weight': 'model-00006-of-00006.safetensors',
 'model.embed_tokens.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.input_layernorm.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.mlp.down_proj.weight': 'model-00001-of-00006.safetensors',
 ...
}

大模型推理

要使用此功能，请确保已安装 Accelerate v0.9.0 和 PyTorch v1.9.0 或更高版本！

from_pretrained() is supercharged with Accelerate’s Big Model Inference feature.

Big Model Inference 会在 PyTorch 的 meta 设备上创建一个 模型骨架。meta 设备不存储任何实际数据，仅存储元数据。

随机初始化的权重仅在加载预训练权重时创建，以避免同时在内存中维护两份模型副本。最大内存占用仅为模型的大小。

有关设备放置的更多信息，请参阅设计设备映射。

Big Model Inference 的第二个特性与模型骨架中权重的加载和分发方式有关。模型权重会分发到所有可用设备上，优先使用最快的设备（通常是 GPU），然后将剩余的权重卸载到较慢的设备（CPU 和硬盘）上。

将这两项功能结合使用，可以降低加载大型预训练模型时的内存占用和加载时间。

将 device_map 设置为 "auto" 以启用大模型推理。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")

你也可以在 device_map 中手动将层分配到设备。它应将模型的所有参数映射到某个设备，但如果整个层都位于同一设备上，则无需详细指定该层中所有子模块的位置。

访问 hf_device_map 属性以查看模型在设备上的分布情况。

device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}
model.hf_device_map

磁盘卸载

当模型过大，无法同时容纳在您的 GPU 和 CPU 内存中时，Big Model Inference 会将剩余的权重卸载到磁盘。只要有任何权重映射到 "disk"，就将 offload_folder 设置为一个目录，因为 Transformers 会将卸载的权重存储在该位置。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    offload_folder="./offload",
)

使用 max_memory 来限制每个设备持有的内存量。Transformers 会按速度顺序填充设备，并将无法容纳的部分发送到磁盘。下例将 GPU 限制为 10GB，CPU 限制为 30GB，因此剩余的权重将保存在 offload_folder 中。

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={0: "10GB", "cpu": "30GB"},
    offload_folder="./offload",
)

将 max_memory={} 设置为指定没有可用的 GPU 或 CPU 内存，并强制将整个模型卸载到磁盘。

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={},
    offload_folder="./offload",
)

磁盘卸载以速度换取内存。在每次前向传播中从磁盘读取权重，比从 CPU 或 GPU 内存读取更慢。完全磁盘卸载需要 accelerate>1.14.0

模型数据类型

dtype 参数用于控制实例化模型权重时所使用的 PyTorch dtype。默认情况下，Transformers 会使用 config.json 中的 dtype 或旧版 torch_dtype 值来加载权重。如果 config.json 中未包含这两个值中的任何一个，Transformers 将使用检查点中第一个浮点权重的 dtype。

通过传入特定的 dtype 来覆盖默认值。

import torch
from transformers import AutoModelForCausalLM

# specific dtype
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it", dtype=torch.float16)

AutoConfig 也接受 dtype 参数，用于从零开始实例化的模型。

import torch
from transformers import AutoConfig, AutoModel

my_config = AutoConfig.from_pretrained("google/gemma-2b", dtype=torch.float16)
model = AutoModel.from_config(my_config)

自定义模型

Custom models builds on Transformers’ configuration and modeling classes, supports the AutoClass API, and are loaded with from_pretrained(). The difference is that the modeling code is not from Transformers.

加载自定义模型时请格外谨慎。虽然 Hub 会为每个仓库提供恶意软件扫描，但你仍应小心操作，以免意外执行恶意代码。

在 from_pretrained() 中设置 trust_remote_code=True 以加载自定义模型。

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

作为额外的安全措施，请从特定的修订版本加载自定义模型，以避免加载可能已更改的模型代码。提交哈希值可以从模型的提交历史中复制。

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

有关更多信息，请参阅自定义模型指南。

Update on GitHub