大模型权重加载之——Meta device

通常对于传统的小参数量模型，我们使用如下步骤完成模型创建，参数加载，转移到指定设备：

#1.创建模型
model = Model('cpu')

#2.在内存中加载其权重（通常称为的对象state_dict）
state_dict = torch.load(checkpoint_path)

#3.在创建的模型中加载这些权重
missing, unexpected = model.load_state_dict(pretrained_dict, strict=False, assign=True)

#4.将模型移动到设备上进行推理
model.to(device)

注意：

• 单GPU修改模型结构 load 参数：model.load_state_dict(pretrained_dict, strict=False, assign=True)中，如果修改了模型结构（增加/删除layers），load预训练权重时，需要设置strict=False, assign=True，对于修改的模型结构层进行随机赋值，因为strict默认是True，这时候就需要严格按照模型中参数的Key值来加载参数，否则将报错Miss Keys。

• 如何加载多GPU训练的模型参数：在执行完函数model = nn.DataParallel(model, device_ids=[0,1,2,3]）这条语句后，会给网络中所有的结构层的名称添加module这个字符，此时，如果我们直接使用 model.load_state_dict(torch.load(“model.pth”),strict=True)将会报错，如果你灵机一动将strict的参数改为False，程序是不会报错了，但是测试结果会低到离谱，因为压根就没有参数加载进来，每一层的名称前都添加了module，所以名称都是不匹配的。 这时候有两种解决问题的方法，一是在加载模型前，依旧使用model = nn.DataParallel (model, device_ids=[0,1,2,3]）给模型每一层名称前添加module的字符。不过当我们想要单卡去测试模型时就遇到问题了，此时我们需要手动删除掉模型名称中的"module."这7个字符，注意是7个，还有个 . 这样做可以自由地更改模型参数的名称，不仅可以删减前缀"module. "，同时也能增加前缀，这个在模型拼接时会比较方便。

import torch
import torch.nn as nn
from collections import OrderedDict
 
model = build_model()
state_dict = torch.load("model.pkl")  # 模型可以保存为pth文件，也可以为pt文件。
# create new OrderedDict that does not contain module.
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k.replace('module.', '')  # remove 'module.'
    new_state_dict[name] = v  # 新字典的key值对应的value为一一对应的值。
# load params
net.load_state_dict(new_state_dict, strict=True)  # 重新加载这个模型。

虽然这种方法在过去几年里效果很好，但非常大的模型使这种方法具有挑战性。假如这里选择的模型有 67 亿个参数。在默认精度下，这意味着仅第 1 步（创建模型）就需要大约26.8GB的​​ RAM（float32 中的 1 个参数占用 4 个字节的内存），这个步骤将消耗大量不必要的开销。

然后，步骤 2 将在内存中加载模型的第二个副本（因此默认精度下 RAM 中还有另外 26.8GB）。如果您尝试加载最大的模型，例如 BLOOM 或 OPT-176B（两者都有 1760 亿个参数），则需要 1.4 TB的 CPU RAM。这有点过分了！而所有这些只是为了在步骤 4 中将模型移动到一个（或多个）GPU 上。

显然，我们需要更智能的东西。在这篇博文中，我们将解释 Accelerate 如何利用 PyTorch 功能加载和运行非常大的模型的推理，即使它们不适合 RAM 或一个 GPU。简而言之，它改变了上述过程，如下所示：

meta device上创建的tensor叫meta tensor，拥有tensor的各种属性，如.size(), .stride(), .requires_grad等，但不包含真实数据。只要您在meta device上，您就可以创建任意大的tensor，而不必担心 CPU（或 GPU）RAM。

# 1.创建一个空的（例如没有权重的）模型
with torch.device("meta"): # metadata是一个没有data的tensor，拥有tensor的各种属性，如.size(), .stride(), .requires_grad等，如果不使用device="meta", 创建model将为所有参数/缓冲区分配内存，造成时间和存储浪费
	model = Flux(configs[name].params)

# 2.在内存中加载部分权重
state_dict = torch.load(checkpoint_path)

# 3.在空模型中加载这些权重
missing, unexpected = model.load_state_dict(pretrained_dict, strict=False, assign=True)

# 4.移动设备上的权重进行推理
model.to(device)

注意：我们不能将Meta Tensor直接转换为 CPU/CUDA Tensor，因为 Meta Tensor 不存储任何真实数据，我们不知道正确的数据值是什么，如下直接使用.to(device)就会报错：

>>> torch.ones(5, device='meta').to(device="cpu")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
NotImplementedError: Cannot copy out of meta tensor; no data!

但我们可以使用.to_empty(device)，将参数和缓冲区移动到指定的设备，而不复制真实数据。

>>> torch.ones(5, device='meta').to_empty(device="cpu")

https://huggingface.co/blog/accelerate-large-models