Hugging Face发布PyTorch新库「Accelerate」：适用于多GPU、TPU、混合精度训练

多数 PyTorch 高级库都支持分布式训练和混合精度训练，但是它们引入的抽象化往往需要用户学习新的 API 来定制训练循环。许多 PyTorch 用户希望完全控制自己的训练循环，但不想编写和维护训练所需的样板代码。Hugging Face 最近发布的新库 Accelerate 解决了这个问题。

「Accelerate」提供了一个简单的 API，将与多 GPU 、 TPU 、 fp16 相关的样板代码抽离了出来，保持其余代码不变。PyTorch 用户无须使用不便控制和调整的抽象类或编写、维护样板代码，就可以直接上手多 GPU 或 TPU。

项目地址：https://github.com/huggingface/accelerate

通过将如下 5 行代码添加到原始的 PyTorch 训练循环中，脚本即可在本地以及任何分布式设置上运行。

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)

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

          optimizer.step()

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

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

          optimizer.step()

除了 API，Accelerate 还提供了一个 CLI 工具，方便启动脚本之前快速配置和测试训练环境，然后启动脚本。指令如下所示：

accelerate config 

accelerate launch my_script.py --args_to_my_script

如果不想自己编写训练循环， PyTorch 之上有许多可以替代 Accelerate 的高级库。

Accelerate 的运作原理

accelerator = Accelerator()

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

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

这是 API 的主体，将准备三种主要类型的对象：models (torch.nn.Module)、optimizers (torch.optim.Optimizer)、dataloaders (torch.data.dataloader.DataLoader)。

模型

模型的准备包括将其包装在适当的容器（例如 DistributedDataParallel）中，然后将其放置在适当的设备上。与普通分布式训练一样，进行保存或访问其特定的方法时，需要先通过 accelerator.unwrap_model（model）解开模型。

优化器

优化器也包装于一个为使混合精度训练运转执行必要操作的特殊容器里。如果状态 dict 是非空的或从检查点加载的，它会准确地处理状态 dict 的设备放置。

数据加载器

此库并不依赖于 DistributedSampler，它实际上可以与传递到数据加载器的采样器一起使用。数据加载器包装于仅在采样器中获取与当前进程相关的索引并将批次放入设备的容器中。

为此，Accelerate 提供了一种实用程序功能来同步更多 RNGs。该功能将在分布式训练期间运行的每个进程上同步随机数生成器。默认情况下，它仅同步采样器的生成器，因此在每个进程中数据扩充都将不同，但是随机改组将是相同的。

accelerator.backward(loss) 

此行代码为向后传递添加了必要的步骤来提高混合精度，但对于其他集成则需要进行一些自定义。

Accelerate 支持的集成包括：

• CPU
• 单 GPU
• 单一节点多 GPU
• 多节点多 GPU
• TPU
• 带有本地 AMP 的 FP16（路线图上的顶点）