Hugging Face Accelerate：使用 PyTorch 运行大型模型

Hugging Face Accelerate 允许用户加载并运行超出普通硬件可用 RAM 或 GPU 内存的大型语言模型（LLM）。通过利用特定的 PyTorch 功能和修改后的加载管线，Accelerate 将模型权重分布到可用的硬件资源上，包括多个 GPU、CPU RAM 和磁盘存储。

克服模型加载中的内存限制

传统的 PyTorch 模型加载遵循线性流程：创建模型、将权重加载到内存中的 state_dict、将这些权重加载到模型中，然后将模型移动到设备上。对于非常大的模型，这一过程在 RAM 方面的开销极其高昂。例如，一个拥有 67 亿参数、使用 float32 精度的模型，仅在初始模型创建时就需要约 26.8GB 的 RAM，加载 state_dict 又需要另外 26.8GB，总计超过 53GB，模型甚至还未进入 GPU。

为了解决此问题，Accelerate 实现了更高效的内存管线：

  1. 创建一个空模型（不带权重）。
  2. 确定设备映射，以决定每一层放置的位置。
  3. 将权重分成小块（shard）加载。
  4. 将这些权重加载到空模型中。
  5. 将权重移动到指定的推理设备上。
  6. 对所有剩余的权重重复上述过程。

通过 Meta 设备实现高效模型初始化

Accelerate 利用 PyTorch 的 “meta” 设备（在 PyTorch 1.9 中引入）来实例化模型，而不为数据分配实际内存。Meta 设备上的张量仅存储形状和数据类型信息，从而可以创建任意大的模型而不消耗 CPU 或 GPU 的 RAM。

由于为 Transformers 库中的每个模型重写以支持 device 参数是不切实际的，Hugging Face 开发了 init_empty_weights() 上下文管理器。它允许任何模型在 meta 设备上实例化为一个“壳”，提供必要的结构信息以计算内存需求，而无需加载实际权重。

自动化设备映射与资源分配

Accelerate 使用 infer_auto_device_map 函数自动将模型权重分配到可用硬件上。系统按以下顺序优先使用资源：GPU、CPU RAM，最后是磁盘 offload。

设备映射配置

根据使用场景，用户可以选择不同的映射策略：

  • "auto" 或 "balanced"：在所有可用 GPU 上均匀分配权重。
  • "balanced_low_0"：在 GPU 上均匀分配权重，但尽量减轻第一块 GPU（GPU 0）的负载，适用于需要将第一块 GPU 用于模型输出的情况（例如文本生成时）。
  • "sequential"：按顺序填充 GPU，可能导致后面的 GPU 未被使用。

为防止系统将单个层拆分到多个设备上（这会破坏残差连接），Accelerate 允许用户指定 no_split_module_classes（例如 ["OPTDecoderLayer"]）。

使用分片检查点降低 RAM 开销

对大多数硬件而言，加载单个巨大的 state_dict 文件是不现实的；例如，BLOOM 模型（1760 亿参数）仅在 bfloat16 精度下加载权重就需要 352GB 的 RAM。为缓解此问题，Hugging Face 使用 分片检查点。

在分片检查点中：

  • pytorch_model.bin.index.json 文件将每个参数名映射到对应的分片文件。
  • 权重被拆分到多个标准的 PyTorch state dict 文件中（例如 BLOOM 有 72 个文件）。

这种方式确保系统只需拥有足够容纳单个最大分片的 RAM（例如 BLOOM 为 7.19GB），而不必容纳整个模型。如果 GPU 和 CPU RAM 不足，Accelerate 可以将权重 offload 到磁盘上指定的 offload_folder。offload_state_dict=True 选项可以在处理其他分片时，临时将驻留在 CPU 的模型部分 offload，从而进一步降低 RAM 使用量。

通过动态 Hook 执行

为了在多个设备上运行被拆分的模型，Accelerate 使用受 PyTorch Hook 启发的机制。dispatch_model 函数会为每个模块及子模块附加在前向传播前后执行的 Hook。这些 Hook 执行以下操作：

  • 确保所有模块的输入与权重位于同一设备上。
  • 在前向传播前立即将权重从 CPU 移动到 GPU 0，前向传播后立即将其返回 CPU。
  • 在前向传播前将权重从磁盘加载到 RAM 再加载到 GPU 0，前向传播后立即释放内存。

虽然此方法是顺序使用 GPU，而非采用复杂的流水线并行，但它使得在显著更小的硬件配置上运行超大模型成为可能。

Sources