将大模型加载到内存

在 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)])

初始化一个参数超过 1000 亿的空模型。在幕后，这依赖于 PyTorch 1.9 中引入的元设备（meta device）。在上下文管理器下进行初始化时，每当创建一个参数，它就会被立即移动到该设备上。

分片检查点

可能你的模型太大，以至于单个副本都无法装入 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"` 的模块不应跨不同设备进行分割。你应该在此处设置所有包含某种残差连接的块。

The 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，它仍然可以进行推理！

设计设备映射

你可以让 Accelerate 处理设备映射的计算，方法是将 `device_map` 设置为支持的选项之一（`"auto"`、`"balanced"`、`"balanced_low_0"`、`"sequential"`），或者如果你想对每个层的放置有更多控制，可以自己创建一个。

当你的 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，并且模型权重使用的 CPU RAM 不超过 30GiB：

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()（或在 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），硬盘卸载可能会非常慢。