🤗 Accelerate 如何通过 PyTorch 运行超大型模型
加载并运行大型模型
Meta AI 和 BigScience 最近开源了超大型语言模型，这些模型无法适应大多数消费级硬件的内存（RAM 或 GPU）。在 Hugging Face，我们使命的一部分是使这些大型模型也易于访问，因此我们开发了工具，即使您没有超级计算机，也可以运行这些模型。本博客文章中选择的所有示例都在免费的 Colab 实例上运行（RAM 和磁盘空间有限），如果您有更多的磁盘空间，请不要犹豫选择更大的检查点。
以下是如何运行 OPT-6.7B
import torch
from transformers import pipeline
# This works on a base Colab instance.
# Pick a larger checkpoint if you have time to wait and enough disk space!
checkpoint = "facebook/opt-6.7b"
generator = pipeline("text-generation", model=checkpoint, device_map="auto", torch_dtype=torch.float16)
# Perform inference
generator("More and more large language models are opensourced so Hugging Face has")
我们稍后会解释每个参数的作用，但首先我们只考虑 PyTorch 中传统的模型加载流程：它通常包括
- 创建模型
- 在内存中加载其权重（通常称为
state_dict的对象） - 将这些权重加载到创建的模型中
- 将模型移动到设备上进行推理
尽管这在过去几年中运行良好，但超大型模型使这种方法面临挑战。这里选择的模型有 67 亿个参数。在默认精度下，这意味着仅第 1 步（创建模型）就需要大约 26.8GB 的 RAM（float32 中的 1 个参数占用 4 字节内存）。这甚至无法适应您在 Colab 上获得的 RAM。
然后，第 2 步将在内存中加载模型的第二个副本（因此默认精度下会再占用 26.8GB RAM）。如果您尝试像这样加载最大的模型，例如 BLOOM 或 OPT-176B（两者都有 1760 亿个参数），您将需要 1.4 TB 的 CPU RAM。这有点过分了！而所有这些只是为了在第 4 步将模型移动到一个（或多个）GPU 上。
显然我们需要更智能的东西。在这篇博客文章中，我们将解释 Accelerate 如何利用 PyTorch 功能加载和运行超大型模型的推理，即使它们不适合 RAM 或一个 GPU。简而言之，它将上述过程更改为这样
- 创建一个空的（例如，没有权重）模型
- 决定每一层将去向何处（当有多个设备可用时）
- 在内存中加载其部分权重
- 将这些权重加载到空模型中
- 将权重移动到设备上进行推理
- 从第 3 步开始重复，加载下一个权重，直到所有权重加载完毕
创建空模型
PyTorch 1.9 引入了一种新的设备，称为 meta 设备。这使我们能够在不附加任何数据的情况下创建张量：meta 设备上的张量只需要一个形状。只要您在 meta 设备上，就可以创建任意大的张量，而无需担心 CPU（或 GPU）RAM。
例如，以下代码将在 Colab 上崩溃
import torch
large_tensor = torch.randn(100000, 100000)
因为这个大型张量需要 4 * 10**10 字节（默认精度是 FP32，所以张量的每个元素占用 4 字节），即 40GB RAM。然而，在 meta 设备上相同的操作可以正常工作
import torch
large_tensor = torch.randn(100000, 100000, device="meta")
如果您尝试显示此张量，PyTorch 将打印以下内容
tensor(..., device='meta', size=(100000, 100000))
正如我们之前所说，这个张量没有关联的数据，只有一个形状。
您可以直接在 meta 设备上实例化模型
large_model = torch.nn.Linear(100000, 100000, device="meta")
但是对于现有模型，这种语法将要求您重写所有建模代码，以便每个子模块都接受并传递一个 device 关键字参数。由于这对于 Transformers 库中的 150 个模型来说不切实际，我们开发了一个上下文管理器，可以为您实例化一个空模型。
以下是实例化 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)
这适用于任何模型，但您会得到一个无法直接使用的 shell：某些操作已针对元设备实现，但并非所有操作都已实现。例如，您可以使用上面定义的 large_model 进行输入，但不能使用 BLOOM 模型。即使使用它，输出也将是元设备的张量，因此您将获得结果的形状，但仅此而已。
作为这项工作的进一步发展，PyTorch 团队正在开发一个新的 FakeTensor 类，它有点像元设备上的张量，但带有设备信息（除了形状和数据类型之外）
由于我们知道每个权重的形状，因此我们可以知道一旦我们完全加载预训练张量，它们将消耗多少内存。因此，我们可以决定如何将模型拆分到 CPU 和 GPU 上。
计算设备映射
在开始加载预训练权重之前，我们需要知道将它们放在何处。这样，我们每次将权重放置在正确位置时，都可以释放 CPU RAM。这可以通过元设备上的空模型完成，因为我们只需要知道每个张量的形状及其数据类型即可计算它将占用多少内存空间。
Accelerate 提供了一个函数，可以从空模型自动确定设备映射。它将尝试最大化所有可用 GPU 的使用，然后是 CPU RAM，最后标记不适合磁盘卸载的权重。让我们使用 OPT-13b 看看。
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)
这将返回一个将模块或权重映射到设备的字典。例如，在具有一个 Titan RTX 的机器上，我们得到以下结果
{'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'}
Accelerate 评估出嵌入层和解码器直到第 9 个块都可以放在 GPU（设备 0）上，然后第 10 个块的一部分需要放在 CPU 上，以及直到第 17 层的所有后续权重。然后第 18 层在 CPU 和磁盘之间拆分，随后的所有层都必须卸载到磁盘
实际使用此设备映射以后将无法工作，因为构成此模型的层具有残差连接（块的输入被添加到块的输出中），因此给定层的所有内容都应该在同一设备上。我们可以通过传递不应使用 no_split_module_classes 关键字参数拆分的模块名称列表来指示 Accelerate
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'}
现在，每个层始终在同一设备上。
在 Transformers 中，当在 from_pretrained() 方法或 pipeline 中使用 device_map 时，这些要保留在同一设备上的块类会自动提供，因此您无需担心它们。请注意，您有以下 device_map 选项（仅当您有多个 GPU 时才相关）
"auto"或"balanced"：Accelerate 将拆分权重，以便每个 GPU 都被同等使用；"balanced_low_0"：Accelerate 会将权重进行拆分，使得每个 GPU 被同等使用，但第一个 GPU 除外，它将尝试使其权重尽可能少（当您希望在某个 GPU 上处理模型输出时，例如使用generate函数时，这会很有用）；"sequential"：Accelerate 将按顺序填充 GPU（因此最后的 GPU 可能完全不会被使用）。
您也可以传入自己的 device_map，只要它遵循我们之前看到的格式（字典层/模块名称到设备）。
最后，请注意，您收到的 device_map 结果取决于所选的数据类型（因为不同类型的浮点数占用不同的空间）。提供 dtype="float16" 将给出不同的结果
device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")
在这种精度下，我们可以将模型安装到 GPU 上的第 21 层
{'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.21': 0,
'model.decoder.layers.22': 'cpu',
...
'model.decoder.layers.37': 'cpu',
'model.decoder.layers.38': 'disk',
'model.decoder.layers.39': 'disk',
'lm_head': 'disk'}
现在我们知道每个权重应该去哪里，我们可以逐步将预训练权重加载到模型中。
分片 state dicts
传统上，PyTorch 模型保存在一个完整的文件中，该文件包含从参数名到权重的映射。这个映射通常被称为 state_dict。以下是 PyTorch 文档中关于保存和加载的摘录：
# 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'))
这对于参数少于 10 亿的模型来说效果很好，但对于更大的模型来说，这在 RAM 中非常耗费资源。BLOOM 模型有 1760 亿个参数；即使将权重保存为 bfloat16 以节省空间，整体仍表示为 352GB。虽然训练该模型的超级计算机可能拥有这么多内存，但要求推理时也拥有这么多内存是不现实的。
这就是为什么 Hugging Face Hub 上的大型模型不是以包含所有权重的一个大文件保存和共享，而是以多个文件保存和共享。例如，如果您访问 BLOOM 模型页面，您会看到有 72 个名为 pytorch_model_xxxxx-of-00072.bin 的文件，每个文件都包含模型权重的一部分。使用这种格式，我们可以将 state dict 的一部分加载到内存中，将权重放入模型中，将它们移动到正确的设备上，然后在进入下一个部分之前丢弃此 state dict 部分。我们不再需要足够的 RAM 来容纳整个模型，而只需要足够的 RAM 来获取最大的检查点部分，我们称之为分片，在 BLOOM 的情况下是 7.19GB。
我们将像 BLOOM 这样保存在多个文件中的检查点称为分片检查点，我们已将其格式标准化如下
- 一个文件（名为
pytorch_model.bin.index.json）包含一些元数据和一个参数名称到文件名的映射，指示在哪里可以找到每个权重 - 所有其他文件都是标准的 PyTorch 状态字典，它们只包含模型的一部分而不是整个模型。您可以在此处查看索引文件的内容。
要将此类分片检查点加载到模型中，我们只需循环遍历各个分片。Accelerate 提供了一个名为 load_checkpoint_in_model 的函数，如果您克隆了 Hub 中的一个仓库，它将为您完成此操作，或者您可以直接使用 Transformers 的 from_pretrained 方法，该方法将为您处理下载和缓存
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
)
在多个设备上运行拆分模型
我们尚未涉及的最后一部分是 Accelerate 如何使您的模型能够在其权重分布在多个 GPU、CPU RAM 和磁盘文件夹中时运行。这通过钩子非常简单地完成。
钩子 是 PyTorch API，它在每次正向传播调用之前执行函数
我们无法直接使用此功能，因为它们只支持正向传播中具有常规参数且没有关键字参数的模型，但我们采用了相同的思想。模型加载后，dispatch_model 函数将向每个模块和子模块添加钩子，这些钩子在每次正向传播之前和之后执行。它们将
- 确保模块的所有输入都在与权重相同的设备上；
- 如果权重已卸载到 CPU，则在正向传播之前将其移动到 GPU 0，并在之后立即移回 CPU；
- 如果权重已卸载到磁盘，则在正向传播之前将其加载到 RAM，然后加载到 GPU 0，并在之后立即释放此内存。
整个过程总结在以下视频中
这样，即使您的 GPU RAM 和 CPU RAM 不足，您的模型也可以加载和运行。您唯一需要的是磁盘空间（以及大量的耐心！）。虽然这个解决方案在您拥有多个 GPU 时相当幼稚（不涉及巧妙的管道并行性，只是按顺序使用 GPU），但它仍然为 BLOOM 带来了相当不错的结果。它允许您在较小的设置上运行模型（尽管速度较慢）。
要了解有关 Accelerate 大型模型推理的更多信息，请参阅文档。
