跳到正文
原文
Hugging Face Blog·· 2022-09-27精选AI 评分71

Accelerate 如何借助 PyTorch 运行超大规模模型

How 🤗 Accelerate runs very large models thanks to PyTorch

AI 导读

Hugging Face 介绍 Accelerate 如何借助 PyTorch 在 RAM 或单卡 GPU 都放不下模型的情况下加载并推理超大语言模型。

推荐理由

拆解了 Accelerate 借助 meta device、sharded checkpoint 和设备映射在有限显存下跑大模型的完整机制。

正文 · AI 翻译

加载并运行大型模型

Meta AI 和 BigScience 最近开源了超大型语言模型，这些模型无法装入大多数消费级硬件的内存（RAM 或 GPU）中。在 HF Mirror，我们的部分使命就是让这些大型模型也能为人所用，因此我们开发了相关工具，让你即使没有超级计算机也能运行这些模型。本文中选取的所有示例都可以在免费的 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 HF Mirror has")

稍后我们会解释每个参数的作用，但首先请看一下 PyTorch 中传统的模型加载流程：它通常包括：

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

虽然这在过去几年里运行得相当不错，但超大型模型让这种方法充满挑战。这里选用的模型有 6.7 十亿 个参数。在默认精度下，这意味着仅第 1 步（创建模型）就会占用大约 26.8GB 的 RAM（float32 格式的 1 个参数在内存中占 4 字节）。这甚至无法放入 Colab 上可用的 RAM 中。

然后第 2 步会在内存中加载模型的第二个副本（因此在默认精度下 RAM 中又会多出 26.8GB）。如果你想尝试加载最大的模型，例如 BLOOM 或 OPT-176B（两者都有 1760 亿个参数），像这样操作，你将需要 1.4 TB 的 CPU RAM。这有点过分了！而这一切只是为了让模型在第 4 步被移到一块（或多块）GPU 上。

显然，我们需要更聪明的做法。在这篇博文中，我们将解释 Accelerate 如何利用 PyTorch 的特性来加载和运行超大模型的推理，即使它们无法装入 RAM 或单块 GPU。简而言之，它将上述流程改为：

  1. 创建一个空模型（例如，不带权重）
  2. 决定每一层将放在哪里（当有多个设备可用时）
  3. 在内存中加载其部分权重
  4. 将这些权重加载到空模型中
  5. 将权重移到设备上进行推理
  6. 从第 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)

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

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

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

计算设备映射

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

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() 方法中使用 device_map，或在 pipeline 中使用时，这些应保留在同一设备上的块类会自动提供，因此你无需担心它们。请注意，对于 device_map（仅当你有多个 GPU 时才相关），你有以下选项：

  • "auto" 或 "balanced"：Accelerate 将拆分权重，使每个 GPU 被均等使用；
  • "balanced_low_0"：Accelerate 将拆分权重，使每个 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 亿的模型来说效果很好，但对于更大的模型，这会占用大量 RAM。BLOOM 模型有 1760 亿个参数；即使权重以 bfloat16 格式保存以节省空间，整体仍占 352GB。虽然训练该模型的超级计算机可能有这么多可用内存，但要求推理时也这样是不现实的。

这就是为什么 HF Mirror Hub 上的大模型不是以包含所有权重的一个大文件来保存和共享的，而是以多个文件的形式。例如，如果你访问 BLOOM 模型页面，你会看到有 72 个名为 pytorch_model_xxxxx-of-00072.bin 的文件，每个文件包含模型权重的一部分。使用这种格式，我们可以在内存中加载状态字典的一部分，将权重放入模型中，把它们移到正确的设备上，然后在继续下一部分之前丢弃这部分状态字典。这样我们不需要有足够容纳整个模型的 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 内存和磁盘文件夹的情况下运行。这通过钩子（hooks）非常简单地实现。

钩子是 PyTorch 的一种 API，它会在每次前向传播之前添加执行的函数。

我们不能直接使用这个功能，因为它们仅支持前向传播中只有常规参数、没有关键字参数的模型，但我们采用了相同的思路。模型加载后，dispatch_model 函数会为每个模块和子模块添加钩子，这些钩子会在每次前向传播前后执行。它们将：

  • 确保模块的所有输入与权重位于同一设备上；
  • 如果权重已被卸载到 CPU，则在正向传递前将其移动到 GPU 0，之后立即移回 CPU；
  • 如果权重已被卸载到磁盘，则在正向传递前将其加载到内存，再加载到 GPU 0，随后立即释放这部分内存。

整个过程总结在以下视频中：

这样，即使你没有足够的 GPU RAM 和 CPU RAM，你的模型也能被加载和运行。你唯一需要的就是磁盘空间（以及大量的耐心！）虽然如果你有多个 GPU，这个方案相当朴素（没有巧妙的流水线并行，只是顺序使用 GPU），但它仍然为 BLOOM 带来了相当不错的结果。而且它能让你在较小的配置上运行模型（尽管速度更慢）。

要了解更多关于 Accelerate 大模型推理的信息，请参阅文档。

来源：Hugging Face Blog · hf-mirror.com

#教程/实践#部署/工程#开源生态#Hugging Face