将大型模型加载到内存

当在 PyTorch 中加载预训练模型时，通常的工作流程如下所示

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

用通俗的英语来说，这些步骤是

1. 使用随机初始化的权重创建模型
2. 从磁盘加载模型权重（通常在名为 state dict 的字典中）
3. 将这些权重加载到模型中

虽然这对于常规大小的模型非常有效，但当我们处理大型模型时，此工作流程存在一些明显的限制：在步骤 1 中，我们在 RAM 中加载模型的完整版本，并花费一些时间随机初始化权重（这将在步骤 3 中被丢弃）。 在步骤 2 中，我们在 RAM 中加载模型的另一个完整版本，其中包含预训练的权重。 如果您要加载一个具有 60 亿参数的模型，这意味着您需要为模型的每个副本提供 24GB 的 RAM，总共需要 48GB（其中一半用于以 FP16 加载模型）。

流程如何运作：快速概述

流程如何运作：使用代码

实例化一个空模型

Accelerate 引入的第一个用于帮助处理大型模型的工具是上下文管理器 init_empty_weights()，它可以帮助您在不使用任何 RAM 的情况下初始化模型，以便可以对任何大小的模型执行步骤 1。 这是它的工作原理

from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)

例如

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

初始化一个参数略多于 100B 的空模型。 在幕后，这依赖于 PyTorch 1.9 中引入的 meta 设备。 在上下文管理器下的初始化期间，每次创建参数时，它都会立即移动到该设备。

分片检查点

您的模型可能非常大，以至于即使单个副本也无法装入 RAM。 这并不意味着它无法加载：如果您有一个或多个 GPU，则有更多内存可用于存储您的模型。 在这种情况下，最好将您的检查点拆分为几个较小的文件，我们称之为检查点分片。

只要您遵循以下格式，Accelerate 就会处理分片检查点：您的检查点应位于文件夹中，其中包含多个包含部分 state dict 的文件，并且应有一个 JSON 格式的索引，其中包含将参数名称映射到包含其权重的文件的字典。 您可以使用 save_model() 轻松分片您的模型。 例如，我们可以有一个包含以下内容的文件夹

first_state_dict.bin
index.json
second_state_dict.bin

其中 index.json 是以下文件

{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}

并且 first_state_dict.bin 包含 "linear1.weight" 和 "linear1.bias" 的权重，second_state_dict.bin 包含 "linear2.weight" 和 "linear2.bias" 的权重

加载权重

Accelerate 引入的第二个工具是一个函数 load_checkpoint_and_dispatch()，它将允许您在空模型中加载检查点。 这支持完整检查点（包含整个 state dict 的单个文件）以及分片检查点。 它还将自动将这些权重分配到您可用的设备（GPU、CPU RAM）上，因此如果您正在加载分片检查点，则最大 RAM 使用量将是最大分片的大小。

如果你想将大型模型推理与 Transformers 模型一起使用，请查看此 文档。

以下是我们如何使用它来加载 GPT2-1.5B 模型。

让我们下载此模型的分片版本。

pip install huggingface_hub

from huggingface_hub import snapshot_download
checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)

为了初始化模型，我们将使用库 minGPT。

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)

然后，使用以下命令加载我们刚刚下载的检查点

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)

通过传递 device_map="auto"，我们告诉 Accelerate 根据可用资源自动确定将模型的每一层放在哪里

• 首先，我们使用 GPU 上可用的最大空间
• 如果我们仍然需要空间，我们将剩余的权重存储在 CPU 上
• 如果 RAM 不足，我们将剩余的权重作为内存映射张量存储在硬盘驱动器上

no_split_module_classes

此参数将指示名称为 "Block" 的某些模块不应跨不同设备拆分。 您应该在此处设置所有包含某种残差连接的块。

device_map

您可以通过访问模型的 hf_device_map 属性来查看 Accelerate 选择的 device_map

model.hf_device_map

{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}

完全可以为您要使用的层创建自己的 device_map，指定要使用的 GPU 设备（数字）、"cpu" 或 "disk" 并将其传入

device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map
)

运行模型

现在我们已经完成了此操作，我们的模型分布在多个设备上，甚至可能在硬盘驱动器上。 但它仍然可以用作常规 PyTorch 模型

from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())

在幕后，Accelerate 向模型添加了钩子，以便

• 在每一层，输入都放在正确的设备上（因此即使您的模型分布在多个 GPU 上，它也能正常工作）
• 对于卸载到 CPU 上的权重，它们会在前向传递之前放在 GPU 上，并在之后立即清除
• 对于卸载到硬盘驱动器上的权重，它们会加载到 RAM 中，然后在前向传递之前放在 GPU 上，并在之后立即清除

这样，即使您的模型不适合其中一个 GPU 或 CPU RAM，它也可以运行以进行推理！

设计 device map

您可以让 Accelerate 通过将 device_map 设置为受支持的选项之一（"auto"、"balanced"、"balanced_low_0"、"sequential"）来处理 device_map 计算，或者如果您想更好地控制每一层的放置位置，则可以自己创建一个。

当您没有足够的 GPU 内存来容纳整个模型时，所有选项都会产生相同的结果（即将所有可以容纳的内容都放在 GPU 上，然后将权重卸载到 CPU 甚至磁盘上，如果 RAM 不足）。

当您拥有的 GPU 内存多于模型大小时，以下是每个选项之间的区别

• "auto" 和 "balanced" 将模型均匀地拆分到所有可用的 GPU 上，使您可以使用大于 1 的批量大小。
• "balanced_low_0" 将模型均匀地拆分到除第一个 GPU 之外的所有 GPU 上，并且仅将不适合其他 GPU 的内容放在 GPU 0 上。 当您需要使用 GPU 0 来处理输出时，此选项非常有用，例如在使用 Transformers 模型的 generate 函数时
• "sequential" 会将它可以容纳的内容放在 GPU 0 上，然后移动到 GPU 1，依此类推（因此如果不需要，则不会使用最后的 GPU）。

首先请注意，您可以使用 max_memory 参数（在 infer_auto_device_map() 以及所有使用它的函数中可用）来限制每个 GPU 上使用的内存。 设置 max_memory 时，您应该传递一个字典，其中包含 GPU 标识符（例如 0、1 等）和 "cpu" 键，用于指定您要用于 CPU 卸载的最大 RAM。 这些值可以是整数（以字节为单位）或表示带有单位的数字的字符串，例如 "10GiB" 或 "10GB"。

这是一个示例，其中我们不想在两个 GPU 中的每一个上使用超过 10GiB，并且不想为模型权重使用超过 30GiB 的 CPU RAM

from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})

此外，如果您对输出执行一些额外的操作而不将它们放回 CPU（例如在 Transformers 的 generate 方法中），并且如果您将输入放在 GPU 上，则该 GPU 将比其他 GPU 消耗更多内存（Accelerate 始终将输出放回输入的设备）。 因此，如果您想优化最大批量大小并且您有许多 GPU，请给第一个 GPU 更少的内存。 例如，在 8x80 A100 设置上的 BLOOM-176B 中，接近理想的映射是

max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}

正如您所见，我们给剩余的 7 个 GPU 的内存比 GPU 0 多 ~50%。

如果您选择完全自己设计 device_map，它应该是一个字典，键是模型的模块名称，值是有效的设备标识符（例如 GPU 的整数）或 "cpu" 用于 CPU 卸载，"disk" 用于磁盘卸载。 键需要覆盖整个模型，然后您可以根据需要定义您的设备映射：例如，如果您的模型有两个块（假设为 block1 和 block2），每个块都包含三个线性层（假设为 linear1、linear2 和 linear3），则有效的设备映射可以是

device_map = {"block1": 0, "block2": 1}

另一个有效的可以是

device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}

另一方面，这个是无效的，因为它没有覆盖模型的每个参数

device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}

仅 CPU 卸载

如果您想在 CPU 上卸载模型，可以使用 cpu_offload()。 这样做的结果是，模型的所有参数都将被卸载，并且只会保留模型状态字典的一个副本。 在前向传播期间，参数将从该状态字典中提取出来，并放置在执行设备上，并在需要时传递，然后再次卸载。

cpu_offload(model, execution_device)

您还可以使用 cpu_offload_with_hook()。 此函数会将模型卸载到 CPU 上，并在执行时将其放回执行设备。 与 cpu_offload() 的区别在于，模型在前向传播后仍保留在执行设备上，并且仅在调用返回的 hook 的 offload 方法时才会再次卸载。 此外，cpu_offload_with_hook() 的性能更高，但内存节省较少。 它适用于在循环中运行模型的管道。

model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

仅磁盘卸载

要执行磁盘卸载，您可以使用 disk_offload()。 这样做的结果是，模型的所有参数都将作为内存映射数组卸载到给定的文件夹中。 在前向传播期间，参数将从该文件夹中访问，并放置在执行设备上，并在需要时传递，然后再次卸载。

disk_offload(model, offload_dir, execution_device)

限制和进一步开发

我们了解当前 API 中的限制

• infer_auto_device_map() （或 device_map="auto" 在 load_checkpoint_and_dispatch() 中）尝试在您执行它时最大化它看到的 GPU 和 CPU RAM。 虽然 PyTorch 非常擅长有效地管理 GPU RAM（并在不需要时将其返回），但这对于 Python 和 CPU RAM 来说并非完全如此。 因此，自动计算的设备映射可能对 CPU 来说过于密集。 如果您由于 RAM 不足而遇到崩溃，请将一些模块移动到磁盘设备。
• infer_auto_device_map() （或 device_map="auto" 在 load_checkpoint_and_dispatch() 中）按顺序分配设备（以避免来回移动），因此如果您的第一层大于您拥有的 GPU 大小，它最终会将所有内容都放在 CPU/磁盘上。
• load_checkpoint_and_dispatch() 和 load_checkpoint_in_model() 目前不对您的状态字典与您的模型的正确性进行任何检查（这将在未来版本中修复），因此如果尝试加载具有不匹配或丢失键的检查点，您可能会遇到一些奇怪的错误。
• 当您的模型在多个 GPU 上拆分时使用的模型并行性是幼稚且未优化的，这意味着一次只有一个 GPU 工作，而其他 GPU 则处于空闲状态。
• 当权重卸载到 CPU/硬盘驱动器上时，没有预取（但我们将在未来版本中为此努力），这意味着权重在需要时而不是之前被放置在 GPU 上。
• 如果您的运行硬件在磁盘和 CPU 之间没有快速通信（如 NVMe），则硬盘驱动器卸载可能会非常慢。