Hugging Face 的 LOGO

Accelerate 文档

快速教程

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

快速入门

启动和运行代码的方式有很多，具体取决于您的训练环境（torchrun、DeepSpeed 等）和可用的硬件。Accelerate 提供了一个统一的接口，用于在不同的分布式设置上启动和训练，让您可以专注于 PyTorch 训练代码，而不是将代码适应这些不同设置的复杂性。这使您可以轻松地扩展您的 PyTorch 代码，以便在 GPU 和 TPU 等硬件的分布式设置上进行训练和推理。Accelerate 还提供了大模型推理功能，使加载和运行通常无法装入内存的大型模型的推理变得更加容易。

本快速入门介绍了 Accelerate 的三个主要功能

  • 用于分布式训练脚本的统一命令行启动接口
  • 一个用于使 PyTorch 训练代码适应不同分布式设置的训练库
  • 大模型推理

统一启动接口

Accelerate 通过由 accelerate config 命令生成的统一配置文件，为任何给定的分布式训练框架（DeepSpeed、FSDP 等）自动选择适当的配置值。您也可以将配置值显式地传递给命令行，这在某些情况下（例如使用 SLURM 时）很有帮助。

但在大多数情况下，您应该始终先运行 accelerate config 来帮助 Accelerate 了解您的训练设置。

accelerate config

accelerate config 命令会在 Accelerate 的缓存文件夹中创建并保存一个 default_config.yaml 文件。该文件存储了您的训练环境配置，帮助 Accelerate 根据您的机器正确启动训练脚本。

在配置完您的环境后，您可以使用 accelerate test 测试您的设置，该命令会启动一个简短的脚本来测试分布式环境。

accelerate test

如果配置文件保存在非默认位置（如缓存），请将 --config_file 添加到 accelerate test 或 accelerate launch 命令中以指定配置文件的位置。

一旦您的环境设置完毕，就可以使用 accelerate launch 启动您的训练脚本！

accelerate launch path_to_script.py --args_for_the_script

要了解更多信息，请查看 启动分布式代码 教程以获取有关启动脚本的更多信息。

我们还有一个 配置示例集，其中展示了许多为各种设置预制的**最小**配置示例。

适配训练代码

Accelerate 的下一个主要功能是 Accelerator 类，它能让您的 PyTorch 代码适应不同的分布式设置。

您只需在训练脚本中添加几行代码，就可以让它在多个 GPU 或 TPU 上运行。

+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ device = accelerator.device
+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for batch in training_dataloader:
      optimizer.zero_grad()
      inputs, targets = batch
-     inputs = inputs.to(device)
-     targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
+     accelerator.backward(loss)
      optimizer.step()
      scheduler.step()
  1. 在训练脚本的开头导入并实例化 Accelerator 类。Accelerator 类会初始化分布式训练所需的一切，并根据代码的启动方式自动检测您的训练环境（例如，单 GPU 的机器、多 GPU 的机器、多台多 GPU 的机器或 TPU）。
from accelerate import Accelerator

accelerator = Accelerator()
  1. 移除模型和输入数据上的 .cuda() 等调用。Accelerator 类会自动将这些对象放置到适当的设备上。

这一步是**可选的**，但被认为是最佳实践，让 Accelerate 来处理设备放置。您也可以在初始化 Accelerator 时传递 device_placement=False 来禁用自动设备放置。如果您想使用 .to(device) 显式地将对象放置在设备上，请确保使用 accelerator.device。例如，如果在将模型放置到 accelerator.device 之前创建优化器，在 TPU 上的训练会失败。

Accelerate 默认情况下不对其自动设备放置使用非阻塞传输，这可能导致不希望的 CUDA 同步。您可以通过在初始化 Accelerator 时，传递一个将 non_blocking=True 设置为 dataloader_config 的 DataLoaderConfiguration 来启用非阻塞传输。像往常一样，非阻塞传输只有在 dataloader 也设置了 pin_memory=True 时才能工作。请注意，从 GPU 到 CPU 使用非阻塞传输可能会导致不正确的结果，如果这导致 CPU 操作在未准备好的张量上执行。

device = accelerator.device
  1. 将所有相关的 PyTorch 训练对象（优化器、模型、数据加载器、学习率调度器）在创建后立即传递给 prepare() 方法。该方法会将模型包装在一个为您的分布式设置优化的容器中，使用 Accelerate 版本的优化器和调度器，并为您的数据加载器创建一个分片版本，以便在 GPU 或 TPU 之间分发。
model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, train_dataloader, lr_scheduler
)
  1. 将 loss.backward() 替换为 backward()，以使用适合您训练设置的正确 backward() 方法。
accelerator.backward(loss)

阅读 Accelerate 的内部机制 指南，了解更多关于 Accelerate 如何适配您的代码的详细信息。

分布式评估

要执行分布式评估，请将您的验证数据加载器传递给 prepare() 方法。

validation_dataloader = accelerator.prepare(validation_dataloader)

分布式设置中的每个设备只接收一部分评估数据，这意味着您应该使用 gather_for_metrics() 方法将您的预测结果汇总在一起。该方法要求每个进程上的所有张量大小相同，因此如果您的张量在每个进程上的大小不同（例如，在一个批次中动态填充到最大长度），您应该使用 pad_across_processes() 方法将您的张量填充到所有进程中的最大尺寸。请注意，张量需要是 1D 的，并且我们沿第一个维度连接张量。

for inputs, targets in validation_dataloader:
    predictions = model(inputs)
    # Gather all predictions and targets
    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
    # Example of use with a *Datasets.Metric*
    metric.add_batch(all_predictions, all_targets)

对于更复杂的情况（例如 2D 张量、不想连接张量、包含 3D 张量的字典），您可以在 gather_for_metrics 中传递 use_gather_object=True。这将返回收集后的对象列表。请注意，在 GPU 张量上使用它支持不佳且效率低下。

数据集末尾的数据可能会被复制，以便批次可以平均分配给所有工作进程。gather_for_metrics() 方法会自动删除重复的数据，以计算更准确的指标。

大模型推理

Accelerate 的大模型推理有两个主要功能，init_empty_weights() 和 load_checkpoint_and_dispatch()，用于加载通常无法装入内存的大型模型进行推理。

请查看处理大模型推理指南，以更好地了解大模型推理的内部工作原理。

空权重初始化

init_empty_weights() 上下文管理器通过创建一个*模型骨架*，并在每次创建参数时将其移动和放置到 PyTorch 的 **meta** 设备上，来初始化任意大小的模型。这样，并非所有权重都立即加载，每次只有一小部分模型加载到内存中。

例如，加载一个空的 Mixtral-8x7B 模型比在 CPU 上完全加载模型和权重所需的内存要少得多。

from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

加载并分派权重

load_checkpoint_and_dispatch() 函数将完整的或分片的检查点加载到空模型中，并自动将权重分布到所有可用的设备上。

device_map 参数决定了每个模型层的放置位置，指定为 "auto" 会首先将它们放置在 GPU 上，然后是 CPU，如果内存仍然不足，则最后作为内存映射张量放置在硬盘上。使用 no_split_module_classes 参数来指示哪些模块不应跨设备拆分（通常是那些具有残差连接的模块）。

from accelerate import load_checkpoint_and_dispatch

model_checkpoint = "your-local-model-folder"
model = load_checkpoint_and_dispatch(
    model, checkpoint=model_checkpoint, device_map="auto", no_split_module_classes=['Block']
)

下一步

既然您已经了解了 Accelerate 的主要功能，您的下一步可以包括

  • 查看教程，以温和地了解 Accelerate。如果您是分布式训练和该库的新手，这将特别有用。
  • 深入研究指南，了解如何在特定用例中使用 Accelerate。
  • 通过阅读概念指南，加深对 Accelerate 内部工作原理的理解。
  • 在API 参考中查找类和命令，以查看可用的参数和选项。
< > 在 GitHub 上更新

←安装 概述→