Hugging Face 发布 Accelerate 库，让 PyTorch 训练脚本适配任意分布式设备

🤗 Accelerate

在任何类型的设备上运行你的原始 PyTorch 训练脚本。

PyTorch 之上的大多数高级库都提供了对分布式训练和混合精度的支持，但它们引入的抽象要求用户如果想自定义底层训练循环，就必须学习一套新的 API。🤗 Accelerate 是为那些喜欢完全掌控自己的训练循环，却又不愿编写（和维护）使用分布式训练（用于单节点或多节点上的多 GPU、TPU 等）或混合精度训练所需的样板代码的 PyTorch 用户而创建的。未来计划包括支持 fairscale、deepspeed、AWS SageMaker 特定的数据并行和模型并行。

它提供两样东西：一个简单且一致的 API，用于抽象那些样板代码；以及一个启动器命令，便于在各种配置上运行这些脚本。

轻松集成！

我们先来看一个例子：

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

+ accelerator = Accelerator()
- device = 'cpu'
+ device = accelerator.device

  model = torch.nn.Transformer().to(device)
  optim = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optim, data = accelerator.prepare(model, optim, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

只需向任何标准 PyTorch 训练脚本添加五行代码，你现在就可以在任何类型的分布式设置下运行该脚本，无论是否使用混合精度。🤗 Accelerate 甚至会为你处理设备放置，因此你可以进一步简化上面的训练循环：

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

+ accelerator = Accelerator()
- device = 'cpu'

- model = torch.nn.Transformer().to(device)
+ model = torch.nn.Transformer()
  optim = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optim, data = accelerator.prepare(model, optim, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
-         source = source.to(device)
-         targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()

相比之下，要让这段代码在分布式训练下运行，需要进行的更改如下：

+ import os
  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from torch.utils.data import DistributedSampler
+ from torch.nn.parallel import DistributedDataParallel

+ local_rank = int(os.environ.get("LOCAL_RANK", -1))
- device = 'cpu'
+ device = device = torch.device("cuda", local_rank)

  model = torch.nn.Transformer().to(device)
+ model = DistributedDataParallel(model)  
  optim = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
+ sampler = DistributedSampler(dataset)
- data = torch.utils.data.DataLoader(dataset, shuffle=True)
+ data = torch.utils.data.DataLoader(dataset, sampler=sampler)

  model.train()
  for epoch in range(10):
+     sampler.set_epoch(epoch)  
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

          loss.backward()

          optimizer.step()

这些更改将使你的训练脚本能在多个 GPU 上运行，但你的脚本随后将无法在 CPU 或单个 GPU 上运行（除非你开始到处添加 if 语句）。更烦人的是，如果你想在 TPU 上测试脚本，就需要修改不同的代码行。混合精度训练也是如此。🤗 Accelerate 的承诺是：

• 将对训练循环的更改保持在最低限度，这样你需要学习的内容尽可能少。
• 让相同的函数适用于任何分布式设置，因此你只需学习一套 API。

它是如何工作的？

为了了解该库在实际中如何工作，我们来看一下需要添加到训练循环中的每一行代码。

accelerator = Accelerator()

除了提供你将使用的主要对象之外，这一行还会从环境中分析分布式训练运行的类型，并执行必要的初始化。你可以通过向此初始化传入 cpu=True 或 fp16=True 来强制在 CPU 上训练或进行混合精度训练。这两个选项也可以通过脚本的启动器来设置。

model, optim, data = accelerator.prepare(model, optim, data)

这是 API 的主要部分，将准备三种主要类型的对象：模型（torch.nn.Module）、优化器（torch.optim.Optimizer）和数据加载器（torch.data.dataloader.DataLoader）。

模型

模型准备包括将其包装在适当的容器中（例如 DistributedDataParallel），并将其放到适当的设备上。与常规分布式训练一样，你需要解包模型以便保存，或访问其特定方法，这可以通过 accelerator.unwrap_model(model) 来完成。

优化器

优化器也被包装在一个特殊容器中，该容器将在 step 中执行必要的操作，以使混合精度正常工作。如果状态字典非空或从检查点加载，它还会正确处理状态字典的设备放置。

数据加载器

这里隐藏着大部分的神奇之处。正如你在代码示例中所见，该库并不依赖于 DistributedSampler，它实际上可以与你传递给 dataloader 的任何采样器配合使用（如果你曾经不得不编写自定义采样器的分布式版本，现在再也不需要了！）。dataloader 被包装在一个容器中，该容器只会获取采样器中与当前进程相关的索引（或者如果你使用 IterableDataset，则跳过其他进程的批次），并将批次放到正确的设备上。

为了实现这一点，Accelerate 提供了一个实用函数，用于在分布式训练期间运行的每个进程上同步随机数生成器。默认情况下，它只同步采样器的 generator，因此每个进程上的数据增强会有所不同，但随机打乱将是相同的。当然，如果需要，你可以使用此实用程序同步更多的 RNG。

accelerator.backward(loss)

最后一行添加了反向传播所需的步骤（主要用于混合精度，但其他集成在这里需要一些自定义行为）。

那评估呢？

评估可以在所有进程上正常运行，或者如果你只想在主进程上运行，可以使用便捷的测试：

if accelerator.is_main_process():
    # Evaluation loop

但你也可以使用 Accelerate 非常轻松地运行分布式评估，以下是你需要在评估循环中添加的内容：

+ eval_dataloader = accelerator.prepare(eval_dataloader)
  predictions, labels = [], []
  for source, targets in eval_dataloader:
      with torch.no_grad():
          output = model(source)

-     predictions.append(output.cpu().numpy())
-     labels.append(targets.cpu().numpy())
+     predictions.append(accelerator.gather(output).cpu().numpy())
+     labels.append(accelerator.gather(targets).cpu().numpy())

  predictions = np.concatenate(predictions)
  labels = np.concatenate(labels)

+ predictions = predictions[:len(eval_dataloader.dataset)]
+ labels = label[:len(eval_dataloader.dataset)]

  metric_compute(predictions, labels)

与训练类似，你需要添加一行来准备你的评估 dataloader。然后你可以使用 accelerator.gather 跨进程收集预测和标签的张量。最后要添加的一行将预测和标签截断为数据集中的示例数量，因为准备好的评估 dataloader 会返回一些额外的元素，以确保每个进程上的批次大小相同。

一个启动器统治一切

使用 Accelerate 的脚本将与你传统的启动器完全兼容，例如 torch.distributed.launch。但记住它们的所有参数有点烦人，而且当你用 4 个 GPU 设置好实例后，你将在大多数训练中使用它们全部。Accelerate 附带了一个方便的 CLI，分两步工作：

accelerate config

这将触发一个关于你设置的小问卷，它会创建一个配置文件，你可以编辑它来为你的训练命令设置所有默认值。然后

accelerate launch path_to_script.py --args_to_the_script

将使用这些默认值启动你的训练脚本。你唯一需要做的就是提供训练脚本所需的所有参数。

为了让这个启动器更加出色，你可以使用它通过 SageMaker 启动一个 AWS 实例。查看 本指南 了解如何操作！

如何参与？

要开始，只需 pip install accelerate 或查看 文档 获取更多安装选项。

Accelerate 是一个完全开源的项目，你可以在 GitHub 上找到它，查看其 文档 或浏览我们的 基本示例。如果你有任何问题或希望库支持的功能，请告诉我们。对于所有问题，论坛 是查看的地方！

对于更复杂的实际示例，你可以查看官方的 Transformers 示例。每个文件夹都包含一个利用 Accelerate 库的 run_task_no_trainer.py！