创建大型模型实例(如何加载大型模型)¶
大型预训练模型的访问障碍之一是其所需的内存量。加载预训练的 PyTorch 模型通常需要以下步骤:
- 创建一个带有随机权重的模型。
- 加载预训练权重。
- 将这些预训练权重放到模型中。
前两个步骤都需要完整的模型副本在内存中,如果模型有几 GB 大小,那么你可能没有足够的内存来存储两个副本。在分布式训练环境中,这个问题更加严重,因为每个进程都会加载预训练模型并在内存中存储两个副本。
为了解决这个问题,_fast_init 参数默认设置为 True,以跳过随机初始化,从而提高加载速度。接下来,本文将介绍如何使用 Transformers 加载大型预训练模型,即使这些模型有较高的内存需求。
分区检查点(Sharded Checkpoints)¶
从 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)
分片检查点的主要优势在于,每次只加载一个分片,从而将内存使用量限制在模型大小和最大分片大小之间。
你也可以使用 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)
分片元数据(Shard Metadata)¶
索引文件决定了哪些键在检查点中以及相应的权重存储在哪里。你可以像加载其他 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())
metadata 键提供了模型的总大小。
index["metadata"]
weight_map 键将每个参数名(通常是 PyTorch 模型中的 state_dict)映射到它所在的分片。
index["weight_map"]
使用 Accelerate 加速大型模型推理¶
确保你已经安装了 Accelerate v0.9.0 或更高版本以及 PyTorch v1.9.0 或更高版本。
从 Transformers v4.20.0 开始,from_pretrained() 方法与 Accelerate 的 Big Model Inference 功能结合使用,可以高效处理非常大的模型!Big Model Inference 在 PyTorch 的 meta 设备上创建一个“模型骨架”,随机初始化的参数仅在加载预训练权重时创建。因此,你不会同时在内存中保留两个模型副本,并且最大内存消耗仅限于完整的模型大小。
要启用 Big Model Inference,可以在 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
# 以下加载方法等价
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"}
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 加载,则需要两倍的内存。
为了节省内存,你可以显式设置 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)