隆重推出 🤗 Accelerate
🤗 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)` 完成。
优化器
优化器也封装在一个特殊容器中，该容器将在步骤中执行必要的操作以使混合精度工作。如果其状态字典非空或从检查点加载，它还将正确处理状态字典的设备放置。
DataLoader
这里隐藏着大部分魔力。正如您在代码示例中看到的那样，该库不依赖于 `DistributedSampler`，它实际上适用于您可能传递给数据加载器的任何采样器（如果您曾经不得不编写自定义采样器的分布式版本，那么就不再需要了！）。数据加载器封装在一个容器中，该容器将只获取采样器中与当前进程相关的索引（或者如果使用 `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)
与训练一样，您需要添加一行来准备您的评估数据加载器。然后您只需使用 `accelerator.gather` 在进程之间收集预测和标签的张量。要添加的最后一行将预测和标签截断为数据集中的示例数量，因为准备好的评估数据加载器将返回更多元素以确保每个进程上的批次大小相同。
一个启动器，统领一切
使用 Accelerate 的脚本将与您的传统启动器完全兼容，例如 `torch.distributed.launch`。但记住所有参数有点烦人，当您设置了 4 个 GPU 的实例时，您将使用所有这些 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`！
