LlamaFactory
算力平台:

将大型模型加载到内存 ​

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

py
import torch

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

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

  1. 使用随机初始化的权重创建模型
  2. 从磁盘加载模型权重（通常是一个称为状态字典的字典）
  3. 将这些权重加载到模型中

虽然这种方法对常规大小的模型非常有效，但在处理大型模型时，这个工作流程有一些明显的局限性：在第一步中，我们在 RAM 中加载了模型的完整版本，并花费一些时间随机初始化权重（这些权重在第三步中会被丢弃）。在第二步中，我们再次在 RAM 中加载了模型的完整版本，这次是带有预训练权重的。如果你加载的是一个有 60 亿参数的模型，这意味着你需要为每个模型副本分配 24GB 的 RAM，总共需要 48GB（其中一半用于以 FP16 格式加载模型）。

过程如何运作：快速概览 ​

过程如何工作：使用代码 ​

实例化一个空模型 ​

Accelerate 介绍的第一个帮助处理大模型的工具是一个上下文管理器 [init_empty_weights]，它可以帮助你在不使用任何 RAM 的情况下初始化模型，从而使得第 1 步可以在任何大小的模型上完成。以下是它的工作原理：

py
from accelerate import init_empty_weights

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

例如：

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

初始化一个参数量略超过1000亿的空模型。在幕后，这依赖于PyTorch 1.9中引入的meta设备。在上下文管理器的作用范围内进行初始化时，每次创建参数时，它都会立即被移动到该设备。

分片检查点 ​

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

只要遵循以下格式，Accelerate 将处理分片检查点：检查点应位于一个文件夹中，包含多个包含部分状态字典的文件，并且应有一个 JSON 格式的索引文件，其中包含一个将参数名称映射到其权重所在文件的字典。你可以使用 [~Accelerator.save_model] 轻松地对模型进行分片。例如，我们可能有一个文件夹，其中包含：

bash
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 内存）上，因此如果你加载的是分片检查点，最大内存使用量将是最大分片的大小。

如果你想使用大型模型推理与 Transformers 模型，请参阅此文档。

以下是如何使用此工具加载 GPT2-1.5B 模型。

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

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

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

bash
git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
py
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)

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

py
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：

py
model.hf_device_map
python
{'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"，并传递这个映射：

python
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 模型一样使用它：

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

设计设备映射 ​

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

当你没有足够的 GPU 内存来容纳整个模型时（即尽可能多地将模型加载到 GPU 上，然后将权重卸载到 CPU 或者在内存不足时卸载到磁盘上），所有选项都会产生相同的结果。

当你有比模型大小更多的 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 上使用的内存不超过 30GiB：

python
max_memory = {
    0: "10GiB",
    1: "10GiB",
    "cpu": "30GiB"
}
python
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 时，接近理想的内存分配方案是：

python
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），一个有效的设备映射可以是：

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

另一个有效的可以是：

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

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

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

仅 CPU 卸载 ​

如果你希望将模型卸载到 CPU，可以使用 [cpu_offload]。这样一来，模型的所有参数都将被卸载，并且只会保留模型状态字典的一个副本。在前向传递过程中，参数将从该状态字典中提取，并放置到执行设备上，按需传递，然后再次卸载。

python
cpu_offload(model, execution_device)

你也可以使用 [cpu_offload_with_hook]。此函数会将模型卸载到 CPU，并在执行时将其放回执行设备。与 [cpu_offload] 的不同之处在于，模型在前向传播后仍保留在执行设备上，只有在调用返回的 hook 的 offload 方法时才会再次卸载。此外，[cpu_offload_with_hook] 的性能更高，但内存节省较少。它适用于在循环中运行模型的管道：

python
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]。这样一来，模型的所有参数将作为内存映射数组卸载到指定的文件夹中。在前向传播过程中，参数将从该文件夹中访问并根据需要放置到传递的执行设备上，然后再次卸载。

python
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 造成过大的负担。如果你因为 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 之间没有快速通信（例如 NVMes），硬盘卸载可能会非常慢。

Copyright © 2024-2025 llamafactory.cn