Transformers 文档

实例化一个大模型

实例化一个大模型

访问非常大的预训练模型的一个障碍是所需的内存量。在加载预训练的PyTorch模型时,通常需要:

  1. 创建一个具有随机权重的模型。
  2. 加载您的预训练权重。
  3. 将这些预训练的权重放入模型中。

前两个步骤都需要在内存中保存模型的完整版本,如果模型的大小达到几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的大模型推理

请确保您已安装Accelerate v0.9.0或更高版本以及PyTorch v1.9.0或更高版本。

从 Transformers v4.20.0 开始,from_pretrained() 方法通过 Accelerate 的 Big Model Inference 功能得到了增强,以高效处理非常大的模型!Big Model Inference 在 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"以使用最优的内存模式加载权重(数据类型会自动从模型权重中推导出来)。

specific dtype
auto dtype
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)
< > Update on GitHub