将大型模型加载到内存中

在 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 格式加载模型）。

此 API 相当新，仍处于实验阶段。虽然我们致力于提供稳定的 API，但未来公开 API 的某些小部分可能会发生变化。

流程如何运作：快速概览

流程如何运作：使用代码

实例化一个空模型

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

initializes 一个拥有略多于 1000 亿参数的空模型。在幕后，这依赖于 PyTorch 1.9 中引入的 meta 设备。在上下文管理器下的初始化过程中，每当创建一个参数时，它会立即被移动到该设备上。

无法直接将在 CPU 或其他设备上以这种方式初始化的模型移动到其他位置，因为它不包含任何数据。此外，使用这种空模型执行前向传播也很可能会失败，因为并非所有操作都支持在 meta 设备上运行。

分片检查点

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

只要遵循以下格式，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"
}

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

加载权重

Accelerate 引入的第二个工具是函数 load_checkpoint_and_dispatch()，它允许你将检查点加载到空模型中。该函数支持完整检查点（包含整个 state dict 的单个文件）以及分片检查点。它还会自动将这些权重分发到你可用的设备（GPU、CPU 内存）上，因此如果你加载的是分片检查点，最大内存占用将等于最大分片的大小。

如果您想使用 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)

然后，使用以下内容加载我们刚刚下载的 checkpoint：

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 内存，它也能运行推理！

此功能仅支持模型的推理，不支持训练。大部分计算都在 torch.no_grad() 上下文管理器中执行，以避免中间激活占用额外的 GPU 内存。

设计设备映射

你可以通过将 device_map 设置为受支持的选项之一（"auto"、"balanced"、"balanced_low_0"、"sequential"）来让 Accelerate 处理设备映射的计算，或者如果你希望对每一层应放置的位置进行更精细的控制，也可以自行创建一个设备映射。

您可以在位于 meta 设备上的模型上推导模型的所有尺寸（从而计算 device_map）。

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

当可用的 GPU 内存大于模型大小时，各选项之间的区别如下：

• "auto" 和 "balanced" 会将模型均匀地拆分到所有可用的 GPU 上，从而允许您使用大于 1 的批处理大小。
• "balanced_low_0" 会将模型均匀地拆分到除第一块 GPU 之外的所有 GPU 上，并且仅在 GPU 0 上放置其他 GPU 无法容纳的部分。当您需要使用 GPU 0 对输出进行一些处理时（例如使用 Transformers 模型的 generate 函数），此选项非常理想。
• "sequential" 会尽可能将模型加载到 GPU 0 上，然后依次加载到 GPU 1 等后续 GPU 上（因此，如果不需要，就不会使用最后的 GPU）。

选项 "auto" 和 "balanced" 目前会产生相同的结果，但如果我们找到更合理的策略，"auto" 的行为未来可能会发生变化，而 "balanced" 将保持稳定。

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

下面是一个示例，其中我们希望每块 GPU 的占用不超过 10GiB，且用于模型权重的 CPU 内存不超过 30GiB：

from accelerate import infer_auto_device_map

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

当 PyTorch 中发生首次内存分配时，它会加载 CUDA 内核，根据 GPU 的不同，大约会占用 1-2GB 的内存。因此，你实际可用的内存总是小于 GPU 的实际容量。要查看实际使用的内存量，请执行 torch.ones(1).cuda() 并查看内存使用情况。

因此，使用 max_memory 创建内存映射时，请确保相应地调整可用内存，以避免内存不足错误。

此外，如果您对输出执行了一些额外操作而没有将其放回 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}

为了达到最高效率，请确保你的设备映射以顺序方式将参数放置在 GPU 上（例如，不要将第一个权重放在 GPU 0 上，然后将权重放在 GPU 1 上，最后又将最后一个权重放回 GPU 0），以避免在 GPU 之间进行大量数据传输。

仅 CPU 卸载

如果要将模型卸载到 CPU，可以使用 cpu_offload()。这样，模型的所有参数都会被卸载，并且只保留一份模型状态字典的副本。在前向传播过程中，参数会从该状态字典中提取出来，放到执行设备上，并在需要时传入，随后再次被卸载。

cpu_offload(model, execution_device)

You can also use cpu_offload_with_hook(). This function will offloads a model on the CPU and puts it back to an execution device when executed. The difference with cpu_offload() is that the model stays on the execution device after the forward and is only offloaded again when the offload method of the returned hook is called. Furthermore, cpu_offload_with_hook() is more performant but less memory saving. It is useful for pipelines running a model in a loop:

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 内存。虽然 PyTorch 非常擅长高效管理 GPU 内存（并在不需要时将其释放），但 Python 和 CPU 内存并非如此。因此，自动计算的设备映射可能对 CPU 的压力过大。如果因内存不足而发生崩溃，请将一些模块移动到磁盘设备上。
• infer_auto_device_map()（或在 load_checkpoint_and_dispatch() 中使用 device_map="auto"）会按顺序分配设备（以避免来回移动数据），因此如果您的第一层比 GPU 的容量更大，最终所有层都会被放到 CPU/磁盘上。
• 目前，load_checkpoint_and_dispatch() 和 load_checkpoint_in_model() 不会检查您的 state dict 与模型之间的一致性（此问题将在未来版本中修复），因此，如果尝试加载键不匹配或缺失的 checkpoint，您可能会遇到一些奇怪的错误。
• 当模型被拆分到多个 GPU 上时，所使用的模型并行是朴素的且未优化的，这意味着在任意给定时刻只有一个 GPU 在工作，其余 GPU 则处于空闲状态。
• 当权重被卸载到 CPU/硬盘时，目前尚不支持预取（我们将在未来版本中对此进行改进），这意味着权重只会在需要时才被加载到 GPU 上，而不是提前加载。
• 如果运行硬件的磁盘与 CPU 之间没有高速通信（例如 NVMe），硬盘卸载可能会非常慢。