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

莫余

已于 2023-10-01 21:50:06 修改

阅读量812

点赞数

CC 4.0 BY-SA版权

分类专栏：多模态文章标签： CUDA python AIGC 大模型 accelerate

于 2023-10-01 21:47:15 首次发布

本文链接：https://blog.csdn.net/qq_44824148/article/details/133466351

多模态专栏收录该内容

14 篇文章

订阅专栏

文章目录

accelerate-large-models

accelerate-large-models

1. 加载和运行大模型

1.1 一般的模型

创建模型
加载权重到内存中
将权重加载到模型中
将模型移到推理设备上

1.2 大模型

创建一个空的（例如无权重）模型
确定每个层的位置（当有多个设备可用时）
在内存中加载其权重的部分
在空模型中加载这些权重
移动设备上的权重以进行推断
对下一部分权重重复步骤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
)