实例化大型模型
访问非常大的预训练模型的一个障碍是所需的内存量。加载预训练的 PyTorch 模型时，通常会
- 创建一个具有随机权重的模型。
- 加载你的预训练权重。
- 将这些预训练权重放入模型中。
前两个步骤都需要模型在内存中的完整版本，如果模型重达几个 GB，则可能没有足够的内存来存储两个副本。在分布式训练环境中，这个问题会加剧，因为每个进程都会加载一个预训练模型并在内存中存储两个副本。
随机创建的模型使用“空”张量进行初始化，这些张量在内存中占用空间，但没有填充它。随机值是在此时此内存块中的任何内容。为了提高加载速度，默认情况下将 _fast_init 参数设置为 True，以跳过所有正确加载的权重的随机初始化。
本指南将向您展示 Transformers 如何帮助您加载大型预训练模型，尽管它们有内存需求。
分片检查点
从 Transformers v4.18.0 开始，大于 10GB 的检查点将由 save_pretrained() 方法自动分片。它被分成几个较小的部分检查点，并创建一个索引文件，将参数名称映射到存储它们的的文件。
最大分片大小由 max_shard_size 参数控制，但默认情况下为 5GB，因为在免费层 GPU 实例上运行更容易，而不会耗尽内存。
例如，让我们分片 BioMistral/BioMistral-7B。
>>> with tempfile.TemporaryDirectory() as tmp_dir:
... model.save_pretrained(tmp_dir, max_shard_size="5GB")
... print(sorted(os.listdir(tmp_dir)))
['config.json', 'generation_config.json', 'model-00001-of-00006.safetensors', 'model-00002-of-00006.safetensors', 'model-00003-of-00006.safetensors', 'model-00004-of-00006.safetensors', 'model-00005-of-00006.safetensors', 'model-00006-of-00006.safetensors', 'model.safetensors.index.json']使用 from_pretrained() 方法重新加载分片检查点。
>>> with tempfile.TemporaryDirectory() as tmp_dir:
... model.save_pretrained(tmp_dir, max_shard_size="5GB")
... new_model = AutoModel.from_pretrained(tmp_dir)对于大型模型，分片检查点的主要优势在于，每个分片在加载上一个分片之后加载，这将内存使用限制在模型大小和最大分片大小。
您还可以直接在模型内部加载分片检查点，而无需使用 from_pretrained() 方法（类似于 PyTorch 的完整检查点的 load_state_dict() 方法）。在这种情况下，请使用 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)分片元数据
索引文件确定检查点中包含哪些键以及相应权重存储在何处。此文件像任何其他 JSON 文件一样加载，您可以从中获取字典。
>>> 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())
dict_keys(['metadata', 'weight_map'])metadata 键提供模型的总大小。
>>> index["metadata"]
{'total_size': 28966928384}weight_map 键将每个参数名称（通常是 PyTorch 模型中的 state_dict）映射到存储它的分片。
>>> 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 或更高版本。
从 Transformers v4.20.0 开始，from_pretrained() 方法通过 Accelerate 的 大型模型推理 功能得到增强，可以有效地处理非常大的模型！大型模型推理在 PyTorch 的 meta 设备上创建了一个模型骨架。预训练权重加载时才会创建随机初始化的参数。这样，您就不会同时在内存中保留模型的两个副本（一个用于随机初始化的模型，一个用于预训练权重），并且消耗的最大内存仅为完整模型的大小。
要在 Transformers 中启用大型模型推理，请在 from_pretrained() 方法中设置 low_cpu_mem_usage=True。
from transformers import AutoModelForCausalLM
gemma = AutoModelForCausalLM.from_pretrained("google/gemma-7b", low_cpu_mem_usage=True)Accelerate 自动将模型权重分发到所有可用的设备上，首先从最快的设备（GPU）开始，然后卸载到较慢的设备（CPU，甚至硬盘）。这可以通过在 from_pretrained() 方法中设置 device_map="auto" 来启用。当您传递 device_map 参数时，low_cpu_mem_usage 会自动设置为 True，因此您无需指定它。
from transformers import AutoModelForCausalLM
# these loading methods are equivalent
gemma = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")
gemma = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto", low_cpu_mem_usage=True)您还可以通过将每一层映射到一个设备来编写自己的 device_map。它应该将所有模型参数映射到一个设备，但是如果整个层都在同一个设备上，则不必详细说明层的所有子模块在哪里。
device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}访问 hf_device_map 属性以查看 Accelerate 如何在设备之间拆分模型。
gemma.hf_device_map
{'model.embed_tokens': 0,
'model.layers.0': 0,
'model.layers.1': 0,
'model.layers.2': 0,
'model.layers.3': 0,
'model.layers.4': 0,
'model.layers.5': 0,
'model.layers.6': 0,
'model.layers.7': 0,
'model.layers.8': 0,
'model.layers.9': 0,
'model.layers.10': 0,
'model.layers.11': 0,
'model.layers.12': 0,
'model.layers.13': 0,
'model.layers.14': 'cpu',
'model.layers.15': 'cpu',
'model.layers.16': 'cpu',
'model.layers.17': 'cpu',
'model.layers.18': 'cpu',
'model.layers.19': 'cpu',
'model.layers.20': 'cpu',
'model.layers.21': 'cpu',
'model.layers.22': 'cpu',
'model.layers.23': 'cpu',
'model.layers.24': 'cpu',
'model.layers.25': 'cpu',
'model.layers.26': 'cpu',
'model.layers.27': 'cpu',
'model.layers.28': 'cpu',
'model.layers.29': 'cpu',
'model.layers.30': 'cpu',
'model.layers.31': 'cpu',
'model.norm': 'cpu',
'lm_head': 'cpu'}模型数据类型
PyTorch 模型权重通常实例化为 torch.float32，如果您尝试以不同的数据类型加载模型，这可能是一个问题。例如，您需要两倍的内存来加载 torch.float32 中的权重，然后再次加载到您所需的数据类型中，例如 torch.float16。
由于 PyTorch 的设计方式，torch_dtype 参数仅支持浮点数据类型。
为了避免浪费内存，请显式地将 torch_dtype 参数设置为所需的数据类型，或者设置 torch_dtype="auto" 以使用最优的内存模式加载权重（数据类型会从模型权重自动推导）。
from transformers import AutoModelForCausalLM
gemma = AutoModelForCausalLM.from_pretrained("google/gemma-7b", torch_dtype=torch.float16)您还可以设置用于从头开始实例化模型的数据类型。
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)