使用🤗 Accelerate简化多GPU/TPU/fp16训练代码

路人与大师
576 阅读4分钟

在深度学习模型的训练过程中，特别是在使用PyTorch进行模型训练时，许多开发者都喜欢手动编写训练循环。然而，当需要在多GPU、TPU或使用混合精度（如fp16）进行训练时，通常需要编写大量的样板代码。这些代码不仅繁琐，而且难以维护。为了简化这一过程，🤗 Accelerate应运而生。

🤗 Accelerate是什么？

🤗 Accelerate 是一个专为PyTorch用户设计的工具。它的目标是将多GPU、TPU以及混合精度相关的样板代码抽象出来，而不改变你的核心代码。这意味着你可以继续使用熟悉的PyTorch语法编写训练循环，而不必担心复杂的分布式训练细节。

示例代码

下面是一个使用🤗 Accelerate的简单示例，展示了如何将标准的PyTorch训练脚本扩展为支持多GPU/TPU和混合精度训练。

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

accelerator = Accelerator()
device = accelerator.device

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

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

model, optimizer, data = accelerator.prepare(model, optimizer, 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)
        optimizer.step()

如你所见，仅需添加五行代码，标准的PyTorch训练脚本就可以在任何单节点或多节点环境中运行，无论是单CPU、单GPU、多GPU还是TPU，甚至可以使用或不使用混合精度（fp8、fp16、bf16）。此外，同样的代码也可以在本地机器上运行，用于调试或训练环境。

🤗 Accelerate还会为你处理设备放置（device placement），这样你可以进一步简化你的训练循环。以下是进一步简化后的代码：

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

accelerator = Accelerator()

model = torch.nn.Transformer()
optimizer = torch.optim.Adam(model.parameters())

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

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

model.train()
for epoch in range(10):
    for source, targets in data:
        optimizer.zero_grad()

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

        accelerator.backward(loss)
        optimizer.step()

通过这种方式，你可以在不同的设备和环境之间无缝切换，而无需修改代码。

🤗 Accelerate 的 CLI 工具

除了核心库之外，🤗 Accelerate 还提供了一个可选的命令行工具（CLI），可以帮助你快速配置和测试训练环境。在你的机器上只需运行：

accelerate config

按照提示回答问题后，工具会自动生成一个配置文件，用于设置默认选项。然后，你可以使用以下命令启动训练脚本：

accelerate launch my_script.py --args_to_my_script

例如，要运行GLUE任务中的MRPC示例，你可以执行：

accelerate launch examples/nlp_example.py

这个CLI工具是可选的，你仍然可以使用python my_script.py或python -m torchrun my_script.py来运行你的脚本。

DeepSpeed 和 MPI 支持

🤗 Accelerate 还支持使用 DeepSpeed 进行训练，以及使用 MPI 进行多CPU训练。对于DeepSpeed，您不需要更改训练代码，所有配置都可以通过accelerate config设置。然而，如果你想在Python脚本中调整DeepSpeed的参数，可以使用DeepSpeedPlugin。

from accelerate import Accelerator, DeepSpeedPlugin

deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
accelerator = Accelerator(mixed_precision='fp16', deepspeed_plugin=deepspeed_plugin)

# 保存模型
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))

如果你想在笔记本环境（如Colab或Kaggle）中运行分布式训练，🤗 Accelerate 还提供了notebook_launcher函数，让你可以在笔记本中启动训练。

为什么选择🤗 Accelerate？

如果你希望在不失去对训练循环的完全控制的情况下，轻松在分布式环境中运行你的训练脚本，🤗 Accelerate 是一个理想的选择。它不是一个高层次的框架，而是一个轻量级的包装器，让你无需学习新的库即可简化多GPU/TPU训练的实现。

何时不该使用🤗 Accelerate？

如果你不想自己编写训练循环，那么🤗 Accelerate可能不适合你。对于这种情况，你可能更倾向于使用PyTorch之上的其他高层库，这些库提供了自动的训练循环管理。

总结

🤗 Accelerate 是一个功能强大且灵活的工具，帮助PyTorch用户简化多设备、多精度的分布式训练。如果你想要快速入门并简化你的分布式训练代码，🤗 Accelerate 无疑是一个值得尝试的选择。

更多详细信息，请查看官方文档。

avatar
文章
阅读
粉丝