将大型模型加载到内存中

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

import torch

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

用简单的英语来说，这些步骤是

1. 使用随机初始化权重创建模型
2. 从磁盘加载模型权重（通常在称为状态字典的字典中）
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)])

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

分片检查点

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

只要您遵循以下格式，Accelerate 就会处理分片检查点：您的检查点应位于一个文件夹中，其中包含几个文件，这些文件包含部分状态字典，并且应该有一个 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()，它允许您在空模型中加载检查点。这支持完整检查点（单个文件包含整个状态字典）以及分片检查点。它还将自动跨您可用的设备（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}

完全可以为您要使用的层创建自己的设备映射，指定要使用的 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设置为支持的选项之一（"auto"、"balanced"、"balanced_low_0"、"sequential"）来让 Accelerate 处理设备映射计算，或者如果您想要更多地控制每一层应该放置的位置，则可以自己创建一个。

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

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

• "auto" 和 "balanced" 将模型均匀地分配到所有可用的 GPU 上，使您能够使用大于 1 的批大小。
• "balanced_low_0" 将模型均匀地分配到除第一个 GPU 之外的所有 GPU 上，并且只将无法放入其他 GPU 的内容放在 GPU 0 上。当您需要使用 GPU 0 来处理输出（例如，在使用 Transformer 模型的generate函数时），此选项非常有用。
• "sequential" 将尽可能多地内容放入 GPU 0，然后移动到 GPU 1，依此类推（因此如果不需要，不会使用最后的 GPU）。

首先要注意，您可以使用max_memory参数限制每个 GPU 上使用的内存（在infer_auto_device_map()和所有使用它的函数中可用）。设置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（例如在 Transformer 的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()（或load_checkpoint_and_dispatch()中的device_map="auto"）在您执行它时尝试最大化它看到的可用的 GPU 和 CPU RAM。虽然 PyTorch 非常擅长有效地管理 GPU RAM（并在不需要时将其释放），但对于 Python 和 CPU RAM 来说并非完全如此。因此，自动计算的设备映射可能会对 CPU 造成过大的压力。如果由于 RAM 不足导致崩溃，请将一些模块移动到磁盘设备。
• infer_auto_device_map()（或load_checkpoint_and_dispatch()中的device_map="auto"）按顺序分配设备（以避免来回移动内容），因此如果您的第一层大于您拥有的 GPU 的大小，它最终会将所有内容都放在 CPU/磁盘上。
• load_checkpoint_and_dispatch()和load_checkpoint_in_model()目前不会对您的状态字典与模型的正确性进行任何检查（这将在未来的版本中修复），因此如果您尝试加载具有不匹配或缺少键的检查点，可能会遇到一些奇怪的错误。
• 当您的模型拆分到多个 GPU 上时使用的模型并行性是简单的，并且没有优化，这意味着在任何给定时间只有一个 GPU 工作，而其他 GPU 处于空闲状态。
• 当权重卸载到 CPU/硬盘上时，没有预取（目前还没有，我们将在未来的版本中进行改进），这意味着权重在需要时才会被放置到 GPU 上，而不是提前放置。
• 如果您的硬件在磁盘和 CPU 之间没有快速的通信（例如 NVMe），则硬盘卸载可能会非常慢。