实例化一个大模型
访问非常大的预训练模型的一个障碍是所需的内存量。在加载预训练的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的大模型推理
请确保您已安装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"
以使用最优的内存模式加载权重(数据类型会自动从模型权重中推导出来)。
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)