Accelerate 文档

快速入门

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始入门

快速入门

根据您的训练环境(torchrunDeepSpeed 等)和可用硬件,有许多方法可以启动和运行您的代码。 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 testaccelerate 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 时传递 DataLoaderConfiguration,并将 non_blocking=True 设置为 dataloader_config 来启用非阻塞传输。 与往常一样,只有当数据加载器也设置了 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 上更新