快速浏览

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

本快速浏览介绍了 Accelerate 的三个主要功能

一个用于分布式训练脚本的统一命令行启动界面
一个训练库，用于将 PyTorch 训练代码改编为在不同的分布式设置上运行
大型模型推理

统一启动界面

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

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

accelerate config

该 accelerate config 命令在 Accelerates 缓存文件夹中创建并保存一个 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()

在训练脚本的开头导入并实例化 Accelerator 类。该 Accelerator 类初始化分布式训练所需的一切，并根据代码的启动方式自动检测您的训练环境（一台带 GPU 的机器、一台带多个 GPU 的机器、多台带多个 GPU 的机器或 TPU 等）。

from accelerate import Accelerator

accelerator = Accelerator()

删除对模型和输入数据调用 .cuda() 的操作。该 Accelerator 类会自动将这些对象放置到适合您的设备上。

此步骤是可选的，但建议您使用 Accelerate 来处理设备放置。您还可以通过在初始化 Accelerator 时传递 device_placement=False 来停用自动设备放置。如果您想使用 .to(device) 将对象显式放置到设备上，请确保使用 accelerator.device。例如，如果您在将模型放置到 accelerator.device 之前创建优化器，则训练在 TPU 上会失败。

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

device = accelerator.device

将所有与 PyTorch 相关的训练对象（优化器、模型、数据加载器、学习率调度器）传递给 prepare() 方法，只要它们创建完毕即可。此方法会将模型包装在针对您的分布式设置优化的容器中，使用 Accelerates 版本的优化器和调度器，并创建数据加载器的分片版本以分布到多个 GPU 或 TPU 上。

model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, train_dataloader, lr_scheduler
)

将 loss.backward() 替换为 backward()，以使用适合您的训练设置的正确 backward() 方法。

accelerator.backward(loss)

阅读 Accelerate 的内部机制指南，以了解有关 Accelerate 如何改编代码的更多详细信息。

分布式评估

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

validation_dataloader = accelerator.prepare(validation_dataloader)

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

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)

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

数据集末尾的数据可能会重复，以便批处理可以在所有工作程序之间平均分配。该 gather_for_metrics() 方法会自动删除重复数据，以计算更准确的指标。

大型模型推理

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

查看处理大型模型进行推理指南，以更好地了解大型模型推理在幕后是如何工作的。

空权重初始化

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

加载和调度权重

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

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

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint="mistralai/Mixtral-8x7B-Instruct-v0.1", device_map="auto", no_split_module_classes=['Block']
)

下一步

现在您已经了解了 Accelerate 的主要功能，您的下一步可能包括

查看教程，逐步了解 Accelerate。如果您是分布式训练和库的新手，这将非常有用。
深入了解指南，了解如何将 Accelerate 用于特定的用例。
通过阅读概念指南，加深您对 Accelerate 内部工作机制的概念理解。
在 API 参考中查找类和命令，了解可用的参数和选项。

< > 在 GitHub 上更新

加速