加载模型

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

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

如果可用，from_pretrained() 方法会加载以 safetensors 文件格式存储的权重。传统上，PyTorch 模型权重使用 pickle 工具序列化，该工具已知不安全。Safetensor 文件更安全且加载速度更快。

from transformers import AutoModelForCausalLM

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

本指南将解释模型是如何加载的、你可以加载模型的不同方式、如何解决非常大模型的内存问题，以及如何加载自定义模型。

模型和配置

所有模型都有一个 configuration.py 文件，其中包含特定属性，例如隐藏层数量、词汇量大小、激活函数等。你还会找到一个 modeling.py 文件，它定义了每个层中发生的层和数学运算。modeling.py 文件采用 configuration.py 中的模型属性并相应地构建模型。此时，你拥有一个具有随机权重的模型，需要对其进行训练才能产生有意义的结果。

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

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

一个简化的模型，例如 AutoModel 或 LlamaModel，它输出隐藏状态。
一个带有特定头部的模型，例如 AutoModelForCausalLM 或 LlamaForCausalLM，用于执行特定任务。

模型类

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

有两种模型类，AutoModel 类和特定于模型的类。

自动模型

特定于模型的类

大型模型

大型预训练模型需要大量内存才能加载。加载过程包括

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

你需要足够的内存来同时容纳模型权重（随机和预训练）的两个副本，这可能因你的硬件而异。在分布式训练环境中，这更具挑战性，因为每个进程都会加载一个预训练模型。

Transformers 通过快速初始化、分片检查点、Accelerate 的大型模型推理功能以及支持低比特数据类型来缓解其中一些与内存相关的挑战。

分片检查点

save_pretrained() 会自动分片大于 50GB 的检查点。这可以减少大型模型的碎片数量并简化文件管理。

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

需要动态权重转换的模型，其内存使用取决于模型大小以及单个转换中最大参数的大小。这通常适用于混合专家（MoE）模型，其中内存使用量是模型大小加上一层上的专家数量。有关模型如何加载的更多信息，请参阅动态权重加载器指南。

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

import json

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() 已通过 Accelerate 的大型模型推理功能得到了增强。

大型模型推理在 PyTorch meta 设备上创建一个模型骨架。meta 设备不存储任何实际数据，只存储元数据。

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

在设计设备映射中了解有关设备放置的更多信息。

大型模型推理的第二个功能与权重如何在模型骨架中加载和分派有关。模型权重被分派到所有可用设备，从最快的设备（通常是 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

模型数据类型

dtype 参数控制模型权重在哪个 PyTorch dtype 中实例化。默认情况下，dtype="auto" 会扫描 config.json 中的 dtype 或旧的 torch_dtype 条目，并以该格式加载权重。如果 config.json 缺少此信息，Transformers 会检查检查点中的第一个浮点权重并采用其数据类型。

通过传递特定的数据类型来覆盖默认值。

import torch
from transformers import AutoModelForCausalLM

# default
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it", dtype="auto")

# 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)

自定义模型

自定义模型建立在 Transformers 的配置和建模类之上，支持 AutoClass API，并使用 from_pretrained() 加载。区别在于建模代码不是来自 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
)

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

在 GitHub 上更新