🤗 加速
在任何设备上运行你的原始 PyTorch 训练脚本。
大多数高级库都支持分布式训练和混合精度，但是它们引入的抽象需要用户学习新的 API，如果他们想要自定义底层的训练循环。🤗 加速是为 PyTorch 用户设计的，他们希望对自己的训练循环拥有完全控制，但又不愿意编写（和维护）使用分布式训练（多个 GPU 在一个或多个节点上，TPU，…）或混合精度训练所需的样板代码。未来的计划包括支持 fairscale、deepseed、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 训练脚本中添加五行代码，现在你可以在任何类型的分布式设置上运行该脚本，以及进行混合精度训练与否。🤗 加速甚至为你处理设备放置，所以你甚至可以进一步简化上面的训练循环：
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 语句）。更令人讨厌的是，如果你想要在 TPUs 上测试你的脚本，你需要更改不同的代码行。混合精度训练也是如此。🤗 加速的承诺是：
- 尽量减少对训练循环的更改，使你需要学习的尽可能少。
- 使同样的函数适用于任何分布式设置，只需学习一个 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) 来完成。
优化器
优化器也被包装在一个特殊的容器中，它将执行必要的操作以使混合精度工作。如果状态字典非空或从检查点加载，则它还会正确处理状态字典的设备放置。
数据加载器
这是大部分魔法隐藏的地方。正如您在代码示例中看到的那样，该库不依赖于 DistributedSampler ，实际上它可以使用您传递给数据加载器的任何采样器来工作（如果您曾经需要编写自定义采样器的分布式版本，现在不再需要！）。数据加载器被包装在一个容器中，该容器仅获取采样器中当前进程相关的索引（或者如果您使用 IterableDataset ，则跳过其他进程的批次）并将批次放置在适当的设备上。
为了使其正常工作，Accelerate 提供了一个实用函数，用于在分布式训练期间同步每个进程上的随机数生成器。默认情况下，它仅同步采样器的 generator ，因此您的数据增强在每个进程上会有所不同，但随机洗牌将保持相同。当然，如果需要，您可以使用此实用程序来同步更多的 RNG。
accelerator.backward(loss)这最后一行添加了反向传播的必要步骤（主要用于混合精度，但其他集成可能需要在此处进行一些自定义行为）。
评估方面怎么样？
评估可以在所有进程上正常运行，或者如果您只希望在主进程上运行评估，您可以使用方便的测试：
if accelerator.is_main_process():
# 评估循环但是您也可以很容易地使用 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示例。每个文件夹都包含一个run_task_no_trainer.py，它利用了Accelerate库！