加载模型

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

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

from transformers import AutoModelForCausalLM

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

本指南解释了模型是如何加载的，加载模型的不同方式，如何克服真正大型模型的内存问题，以及如何加载自定义模型。

模型和配置

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

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

1. 一个最基本的模型，如 AutoModel 或 LlamaModel，输出隐藏状态。
2. 一个附加了特定头部的模型，如 AutoModelForCausalLM 或 LlamaForCausalLM，用于执行特定任务。

对于每种模型类型，每种机器学习框架（PyTorch、TensorFlow、Flax）都有一个单独的类。选择您正在使用的框架的相应前缀。

from transformers import AutoModelForCausalLM, MistralForCausalLM

# load with AutoClass or model-specific class
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", , torch_dtype="auto", device_map="auto")
model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", , torch_dtype="auto", device_map="auto")

模型类

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

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

from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModelForQuestionAnswering

# use the same API for 3 different tasks
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForSequenceClassification.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForQuestionAnswering.from_pretrained("meta-llama/Llama-2-7b-hf")

from transformers import AutoModelForCausalLM

# use the same API to load 3 different models
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("google/gemma-7b")

大型模型

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

1. 创建具有随机权重的模型
2. 加载预训练权重
3. 将预训练权重放置在模型上

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

Transformers 通过快速初始化、分片检查点、Accelerate 的 Big Model Inference 功能以及支持较低位数据类型，减少了这些与内存相关的挑战。

快速初始化

PyTorch 模型使用随机权重或“空”张量实例化，这些张量占用内存空间而不填充它。

如果预训练权重已正确初始化，Transformers 通过使用 _fast_init 参数跳过随机权重初始化，从而提高了加载速度。默认情况下，此参数设置为 True。

分片检查点

大于 10GB 的检查点，save_pretrained() 方法会自动对其进行分片。

每个分片在加载完前一个分片后按顺序加载，从而将内存使用量限制为仅模型大小和最大分片大小。

max_shard_size 参数默认为每个分片 5GB，因为这样更容易在免费层 GPU 实例上运行，而不会耗尽内存。

例如，在 save_pretrained() 中为 BioMistral/BioMistral-7B 创建一些分片检查点。

from transformers import AutoModel
import tempfile
import os

model = AutoModel.from_pretrained("biomistral/biomistral-7b")
with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="5GB")
    print(sorted(os.listdir(tmp_dir)))

使用 from_pretrained() 重新加载分片检查点。

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir)
    new_model = AutoModel.from_pretrained(tmp_dir)

分片检查点也可以使用 load_sharded_checkpoint() 直接加载。

from transformers.modeling_utils import load_sharded_checkpoint

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="5GB")
    load_sharded_checkpoint(model, tmp_dir)

save_pretrained() 方法创建一个索引文件，该文件将参数名称映射到它们存储的文件。索引文件有两个键：metadata 和 weight_map。

import json

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="5GB")
    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',
 ...
}

大模型推理

from_pretrained() 通过 Accelerate 的 Big Model Inference 功能进行了增强。

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

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

Big Model Inference 的第二个功能与权重如何在模型骨架中加载和分派有关。模型权重跨所有可用设备分派，首先从最快的设备（通常是 GPU）开始，然后将任何剩余权重卸载到较慢的设备（CPU 和硬盘）。

这两个功能结合起来，减少了大型预训练模型的内存使用量和加载时间。

将 device_map 设置为 "auto" 以启用 Big Model Inference。这也会将 low_cpu_mem_usage 参数设置为 True，这样 CPU 内存中使用的内存不会超过模型大小的 1 倍。

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

模型数据类型

默认情况下，PyTorch 模型权重以 torch.float32 初始化。以不同的数据类型（如 torch.float16）加载模型需要额外的内存，因为模型会以所需的数据类型再次加载。

显式设置 torch_dtype 参数以直接以所需的数据类型初始化模型，而不是两次加载权重（torch.float32，然后是 torch.float16）。您也可以设置 torch_dtype="auto" 以自动以权重存储的数据类型加载权重。

from transformers import AutoModelForCausalLM

gemma = AutoModelForCausalLM.from_pretrained("google/gemma-7b", torch_dtype=torch.float16)

对于从头开始实例化的模型，也可以在 AutoConfig 中配置 torch_dtype 参数。

import torch
from transformers import AutoConfig, AutoModel

my_config = AutoConfig.from_pretrained("google/gemma-2b", torch_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
)

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