【Accelerate】accelerate-large-models （RuntimeError: Expected all tensors to be on the same device……）

accelerate-large-models

1. 加载和运行大模型

1.1 一般的模型

1. 创建模型
2. 加载权重到内存中
3. 将权重加载到模型中
4. 将模型移到推理设备上

1.2 大模型

1. 创建一个空的（例如无权重）模型
2. 确定每个层的位置（当有多个设备可用时）
3. 在内存中加载其权重的部分
4. 在空模型中加载这些权重
5. 移动设备上的权重以进行推断
6. 对下一部分权重重复步骤3，直到所有权重都被加载

2. 创建一个空模型

比如说，这个是无法创建的（默认精度FP32）：

import torch

large_tensor = torch.randn(100000, 100000)

而在meta device上是可以工作的：

import torch

large_tensor = torch.randn(100000, 100000, device="meta")
# tensor(..., device='meta', size=(100000, 100000))

在实际过程中，我们无法对每一个张量的device属性进行修改。

所以，我们一般这样实例化（以BLOOM为例）：

from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("bigscience/bloom")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

3. 计算设备映射

from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

device_map = infer_auto_device_map(model)

运行返回的结果如下：

{'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.9': 0,
 'model.decoder.layers.10.self_attn': 0,
 'model.decoder.layers.10.activation_fn': 0,
 'model.decoder.layers.10.self_attn_layer_norm': 0,
 'model.decoder.layers.10.fc1': 'cpu',
 'model.decoder.layers.10.fc2': 'cpu',
 'model.decoder.layers.10.final_layer_norm': 'cpu',
 'model.decoder.layers.11': 'cpu',
 ...
 'model.decoder.layers.17': 'cpu',
 'model.decoder.layers.18.self_attn': 'cpu',
 'model.decoder.layers.18.activation_fn': 'cpu',
 'model.decoder.layers.18.self_attn_layer_norm': 'cpu',
 'model.decoder.layers.18.fc1': 'disk',
 'model.decoder.layers.18.fc2': 'disk',
 'model.decoder.layers.18.final_layer_norm': 'disk',
 'model.decoder.layers.19': 'disk',
 ...
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

从结果中我们可以看出：

• 从第0层到第9层位于GPU 0上。
• 第10层的前一部分在GPU 0 上，后一部分在CPU上。
• 第11层到第17层位于CPU上。
• 第18层的前一部分在CPU 上，后一部分在磁盘上。
• 第19层到第39层位于磁盘上。

由于每一层需要在同一设备上，才是可行的。

因此，应该添加：

device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"])

代码如下：

from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)
device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"])

这将返回：

'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.9': 0,
 'model.decoder.layers.10': 'cpu',
 'model.decoder.layers.11': 'cpu',
 ...
 'model.decoder.layers.17': 'cpu',
 'model.decoder.layers.18': 'disk',
 ...
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

device_map 模式的解释：

• “auto” 或 “balanced”：Accelerate 将分割权重，以确保每个 GPU 的负载相等。
• “balanced_low_0”：Accelerate 将分割权重，以确保每个 GPU 的负载相等，但第一个 GPU 会尽量保持尽量少的权重（当您希望在一个 GPU 上处理模型的输出时很有用，比如使用生成函数时）。
• “sequential”：Accelerate 将按顺序填充 GPU（因此最后的 GPU 可能根本不会被使用）。

4. 状态分层

4.1 传统保存/加载权重

# Save the model weights
torch.save(my_model.state_dict(), 'model_weights.pth')

# Reload them
new_model = ModelClass()
new_model.load_state_dict(torch.load('model_weights.pth'))

4.2 large-models

Hugging Face Hub 上的大型模型不是通过包含所有权重的一个大文件来保存和共享，而是使用其中的几个权重来保存和共享。

您进入BLOOM 模型页面，您将看到有 72 个名为 的文件pytorch_model_xxxxx-of-00072.bin（每个大约为 7.19GB。），每个文件都包含部分模型权重。使用这种格式，我们可以将状态字典的一部分加载到内存中，将权重放入模型中，将它们移动到正确的设备上，然后在进入下一个之前丢弃此部分。

import torch
from transformers import AutoModelForCausalLM

# Will error
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float16)

如果没有足够的 GPU 和 CPU RAM，会收到一条错误消息，指示您需要传递一个文件夹，在该文件夹中将卸载应存储在磁盘上的权重。

错误信息如下：

ValueError: The current `device_map` had weights offloaded to the disk. Please provide an 
`offload_folder` for them.

解决办法：

import torch
from transformers import AutoModelForCausalLM

# Will go out of RAM on Colab
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map="auto", offload_folder="offload", torch_dtype=torch.float16
)

如果尝试加载一个非常大的模型，除了 CPU 卸载之外还需要一些磁盘卸载，当加载检查点的最后几个分片时，您可能会耗尽 RAM，因为模型的某部分仍然驻留在 CPU 上并占用空间。如果是这种情况，请使用选项 offload_state_dict=True，在加载所有权重后，临时卸载留在 CPU 上的模型部分，然后在处理所有权重后重新加载到 RAM 中。

import torch
from transformers import AutoModelForCausalLM

checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map="auto", offload_folder="offload", offload_state_dict = True, torch_dtype=torch.float16
)

这将适应Colab，但将非常接近使用所有可用的RAM，因此在尝试生成预测时可能会耗尽RAM。为了获得我们可以使用的模型，我们需要在磁盘上卸载一层。我们可以通过获取在前一节中计算的 device_map，稍微调整它，然后将其传递给 from_pretrained 调用来实现这一点。

import torch
from transformers import AutoModelForCausalLM

checkpoint = "facebook/opt-13b"
device_map["model.decoder.layers.37"] = "disk"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map=device_map, offload_folder="offload", offload_state_dict = True, torch_dtype=torch.float16
)

5. 在多个设备上运行模型拆分

hooks是一个 PyTorch API，它添加在每次转发调用之前执行的函数

我们无法直接使用它，因为它们仅支持具有常规参数的模型，并且在前向传递中不支持关键字参数，但我们采用了相同的想法。加载模型后，该dispatch_model函数将向每个前向传递之前和之后执行的每个模块和子模块添加hooks。他们将：

• 确保模块的所有输入与权重位于同一设备上；
• 如果权重已卸载到 CPU，则在前向传递之前将它们移至 GPU 0，并在之后返回到 CPU；
• 如果权重已卸载到磁盘，则在前向传递之前将它们加载到 RAM 中，然后加载到 GPU 0 上，并在之后释放该内存。

6. 总结

视频讲解：https://www.youtube.com/watch?v=MWCSGj9jEAo

此方法需要预先估计，每一层一定是在同一个设备上的。

参考：https://github.com/huggingface/blog/blob/main/accelerate-large-models.md