跳到正文
北京时间
原文
Hugging Face：Blog·· 2022-09-27精选AI 评分71

Hugging Face 讲解 Accelerate 如何借助 PyTorch 加载和运行超大模型

How 🤗 Accelerate runs very large models thanks to PyTorch

AI 导读

Hugging Face 发布博客讲解 Accelerate 如何加载和运行超出单机内存的超大模型。核心做法是用 PyTorch meta device 创建空模型、用 infer_auto_device_map 自动计算 GPU、CPU 与磁盘的设备分配，并按分片加载权重，通过 hooks 在每次前向传播前后搬运权重。

推荐理由

原文逐层拆解 Accelerate 借助 PyTorch meta device 与分片加载运行超大模型的机制，并给出可在免费 Colab 复现的做法。

正文 · AI 翻译

加载并运行大型模型

Meta AI 和 BigScience 最近开源了非常大的语言模型，这些模型无法放入大多数消费级硬件的内存（RAM 或 GPU）中。在 Hugging Face，我们的部分使命是让即使这些大型模型也能被访问，因此我们开发了工具，让您即使没有超级计算机也能运行这些模型。本博客文章中选取的所有示例都可以在免费的 Colab 实例上运行（内存和磁盘空间有限），如果您有更多磁盘空间，请不要犹豫，选择更大的检查点。

以下是我们如何运行 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 中传统的模型加载流程：它通常包括：

  1. 创建模型
  2. 在内存中加载其权重（通常在一个名为 state_dict 的对象中）
  3. 将这些权重加载到创建的模型中
  4. 将模型移动到设备上进行推理

虽然这在过去几年中效果很好，但非常大的模型使这种方法具有挑战性。这里选择的模型有 6.7 十亿 个参数。在默认精度下，这意味着仅步骤 1（创建模型）就会占用大约 26.8GB 的 RAM（一个 float32 参数在内存中占 4 字节）。这甚至无法放入您在 Colab 上获得的内存中。

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

显然我们需要更智能的方法。在这篇博客文章中，我们将解释 Accelerate 如何利用 PyTorch 的功能来加载和运行非常大的模型的推理，即使它们无法放入 RAM 或一个 GPU 中。简而言之，它像这样改变了上述过程：

  1. 创建一个空模型（例如，没有权重）
  2. 决定每一层将放在哪里（当有多个设备可用时）
  3. 在内存中加载其部分权重
  4. 将这些权重加载到空模型中
  5. 将权重移动到设备上进行推理
  6. 从步骤 3 重复下一个权重，直到所有权重都加载完毕

创建一个空模型

PyTorch 1.9 引入了一种新的设备类型，称为 meta 设备。这允许我们创建没有任何数据附加的 tensor：meta 设备上的 tensor 只需要一个形状。因此，只要您在 meta 设备上，就可以创建任意大的 tensor，而无需担心 CPU（或 GPU）RAM。

例如，以下代码在 Colab 上会崩溃：

import torch

large_tensor = torch.randn(100000, 100000)

因为这个大 tensor 需要 4 * 10**10 字节（默认精度是 FP32，所以 tensor 的每个元素占 4 字节），因此需要 40GB 的 RAM。然而，在 meta 设备上做同样的事情却完全没问题：

import torch

large_tensor = torch.randn(100000, 100000, device="meta")

如果您尝试显示这个 tensor，PyTorch 将打印如下内容：

tensor(..., device='meta', size=(100000, 100000))

正如我们之前所说，这个 tensor 没有关联数据，只有一个形状。

您可以直接在 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)

这适用于任何模型，但你会得到一个无法直接使用的空壳：有些操作已在 meta device 上实现，但尚未全部实现。例如，这里你可以使用上面定义的 large_model 并传入输入，但不能使用 BLOOM 模型。即使使用它，输出也将是 meta device 上的张量，因此你只能得到结果的形状，仅此而已。

作为进一步的工作，PyTorch 团队正在开发一个新的 类 FakeTensor，它有点像 meta device 上的张量，但除了形状和 dtype 之外，还带有设备信息。

既然我们知道每个权重的形状，那么一旦完全加载预训练张量，我们就能知道它们将消耗多少内存。因此，我们可以决定如何将模型拆分到 CPU 和 GPU 上。

计算设备映射

在开始加载预训练权重之前，我们需要知道要将它们放在哪里。这样，每当我们将一个权重放到正确的位置后，就可以释放 CPU 内存。这可以在 meta device 上的空模型上完成，因为我们只需要知道每个张量的形状及其 dtype 来计算它将占用多少内存。

Accelerate 提供了一个函数，可以从空模型自动确定设备映射。它会尝试最大化利用所有可用的 GPU，然后是 CPU 内存，最后将不适合的权重标记为磁盘卸载。让我们用 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 上的权重尽可能少（例如，当你想要在一个 GPU 上处理模型的输出时很有用，比如使用 generate 函数时）；
  • "sequential"：Accelerate 将按顺序填充 GPU（因此最后几个 GPU 可能完全不被使用）。

你也可以传入自己的 device_map，只要它遵循我们之前看到的格式（将层/模块名称映射到设备的字典）。

最后，请注意，你收到的 device_map 结果取决于所选的 dtype（因为不同类型的浮点数占用不同的空间）。提供 dtype="float16" 会得到不同的结果：

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

在这种精度下，我们可以将模型直到第 21 层都放在 GPU 上：



{'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'}

现在我们知道了每个权重应该放在哪里，就可以逐步将预训练权重加载到模型中了。

分片状态字典

传统上，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 亿的模型来说效果很好，但对于更大的模型，这会非常消耗内存。BLOOM 模型有 1760 亿个参数；即使权重以 bfloat16 保存以节省空间，整体仍占 352GB。虽然训练这个模型的超级计算机可能有这么多可用内存，但要求推理时也这样是不现实的。

这就是为什么 Hugging Face Hub 上的大模型不是保存和共享为一个包含所有权重的大文件，而是多个文件。例如，如果你访问 BLOOM 模型页面，你会看到有 72 个名为 pytorch_model_xxxxx-of-00072.bin 的文件，每个文件包含模型权重的一部分。使用这种格式，我们可以在内存中加载状态字典的一部分，将权重放入模型，将它们移动到正确的设备上，然后在继续下一个之前丢弃这个状态字典部分。我们不需要有足够的内存来容纳整个模型，只需要足够的内存来获取最大的检查点部分，我们称之为分片，对于 BLOOM 来说就是 7.19GB。

我们将像 BLOOM 这样保存在多个文件中的检查点称为分片检查点，并且我们已经将其格式标准化如下：

  • 一个文件（称为 pytorch_model.bin.index.json）包含一些元数据和参数名到文件名的映射，指示在哪里找到每个权重
  • 所有其他文件都是标准的 PyTorch 状态字典，它们只包含模型的一部分而不是整个模型。你可以在这里查看索引文件的内容。

要将这样的分片检查点加载到模型中，我们只需要遍历各个分片。如果你已经克隆了 Hub 上的某个仓库，Accelerate 提供了一个名为 load_checkpoint_in_model 的函数来为你完成这个操作，或者你可以直接使用 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 内存不足而被卸载到磁盘上，你会收到一个错误，指示你需要传递一个文件夹，用于存放应该存储在磁盘上的权重：

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 卸载的基础上再进行磁盘卸载，那么在加载检查点的最后几个分片时可能会耗尽内存，因为留在 CPU 上的那部分模型占用了空间。如果是这种情况，请使用 offload_state_dict=True 选项，在加载所有权重时暂时卸载留在 CPU 上的那部分模型，并在所有权重处理完毕后将其重新加载到内存中

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 中运行，但会非常接近使用所有可用内存，以至于当你尝试生成预测时会耗尽内存。为了得到一个我们可以使用的模型，我们需要再卸载一层到磁盘上。我们可以通过获取上一节中计算的 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 大模型推理的信息，请参阅 文档。

来源：Hugging Face：Blog · huggingface.co

#教程/实践#部署/工程#Hugging Face