Press "Enter" to skip to content

如何使用🤗 Accelerate和PyTorch运行非常大的模型

Published July 15, 2023 by 四海吧

加载和运行大型模型

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)

# 执行推断
generator("越来越多的大型语言模型被开源，所以Hugging Face有")

稍后我们将解释这些参数的作用，但首先只需考虑PyTorch中传统的模型加载流程：通常包括以下步骤：

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

虽然在过去的几年中，这种方法运行得很好，但是非常大的模型使得这种方法变得具有挑战性。在此例中，选择的模型具有67亿个参数。在默认精度下，这意味着仅第1步（创建模型）将占用大约26.8GB的RAM（float32格式的1个参数在内存中占用4个字节）。这甚至无法适应Colab上的RAM。

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

显然，我们需要一些更智能的方法。在本博客文章中，我们将解释如何利用PyTorch的功能来加载和运行非常大的模型，即使它们不适合RAM或单个GPU。简而言之，它将上述过程更改如下：

  1. 创建一个空的模型（例如，没有权重）
  2. 确定每个层将要放置的位置（当有多个设备可用时）
  3. 将部分权重加载到内存中
  4. 将这些权重加载到空模型中
  5. 将权重移动到用于推断的设备上
  6. 重复从第3步开始的过程，直到加载所有权重

创建一个空模型

PyTorch 1.9引入了一种称为元设备（meta device）的新设备类型。这使我们能够创建没有任何数据附加到其中的张量：在元设备上的张量只需要一个形状即可。只要您在元设备上，您就可以创建任意大的张量，而无需担心CPU（或GPU）的RAM。

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

import torch

large_tensor = torch.randn(100000, 100000)

因为这个大张量需要4 * 10**10字节（默认精度为FP32，所以张量的每个元素占用4个字节），即40GB的RAM。然而，在元设备上执行相同的操作完全没有问题：

import torch

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

如果尝试显示此张量，PyTorch将打印如下：

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

正如我们之前所说，此张量没有与之关联的数据，只有一个形状。

您可以直接在元设备上实例化模型：

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)

这适用于任何模型，但您会得到一个无法直接使用的外壳：一些操作已经针对元设备实现，但不是全部。例如，您可以使用上面定义的large_model进行输入，但不能使用 BLOOM 模型。即使使用它，输出也将是元设备的张量，因此您将得到结果的形状，但不会得到更多信息。

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

由于我们知道每个权重的形状，我们可以知道在完全加载预训练张量后它们将占用多少内存。因此，我们可以决定如何在 CPU 和 GPU 之间划分我们的模型。

计算设备映射

在开始加载预训练权重之前，我们需要知道我们想要将它们放在哪里。这样我们就可以在将权重放在正确位置后释放 CPU 内存。这可以通过在元设备上使用空模型来实现，因为我们只需要知道每个张量的形状和数据类型，就可以计算它将占用多少内存空间。

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 和磁盘之间进行划分，而后续层必须全部卸载到磁盘上

后续使用这个设备映射将无法工作，因为构成这个模型的层具有残差连接（块的输入与块的输出相加），所以给定层的所有组成部分都应该在同一设备上。我们可以通过传递一个模块名称列表来指示给 Accelerate，这些模块不应该被划分，使用no_split_module_classes关键字参数：

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（因此最后几个可能根本不会被使用）。

您还可以传递自己的 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 文档中关于保存和加载的摘录：

# 保存模型权重
torch.save(my_model.state_dict(), 'model_weights.pth')

# 加载模型权重
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 的文件，每个文件包含模型权重的一部分。使用这种格式，我们可以将状态字典的一部分加载到内存中，将权重放入模型中，将其移动到正确的设备上，然后在进入下一个状态字典部分之前丢弃该状态字典部分。与需要足够的 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和磁盘文件夹上运行。这是通过使用hooks来实现的。

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

由于hooks仅支持具有常规参数且在其forward传递中没有关键字参数的模型，因此我们无法直接使用它们，但我们采用了相同的思路。一旦加载了模型，dispatch_model函数将向每个模块和子模块添加在每次forward传递之前和之后执行的hooks。它们将：

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

整个过程在下面的视频中进行了总结：

通过这种方式，即使您没有足够的GPU RAM和CPU RAM，也可以加载和运行您的模型。您唯一需要的是磁盘空间（和大量的耐心！）虽然如果您有多个GPU，这种解决方案可能相对简单（没有使用聪明的管道并行性，只是依次使用GPU），但对于BLOOM来说，它仍然可以产生相当不错的结果。它还允许您在较小的设置上运行模型（尽管速度较慢）。

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

Published in 四海

Leave a Reply

Your email address will not be published. Required fields are marked *