使用 🤗 Accelerate 进行分布式训练
随着模型变得越来越大，并行化已成为一种策略，用于在有限的硬件上训练更大的模型，并将训练速度提高几个数量级。在 Hugging Face，我们创建了 🤗 Accelerate 库来帮助用户轻松地在任何类型的分布式设置上训练 🤗 Transformers 模型，无论是同一台机器上的多个 GPU 还是多台机器上的多个 GPU。在本教程中，学习如何自定义你的原生 PyTorch 训练循环以在分布式环境中启用训练。
设置
开始使用，安装 🤗 Accelerate
pip install accelerate
然后导入并创建一个 Accelerator 对象。 Accelerator 会自动检测你的分布式设置类型并初始化训练所需的所有组件。你不需要显式地将你的模型放在设备上。
>>> from accelerate import Accelerator
>>> accelerator = Accelerator()准备加速
下一步是将所有相关的训练对象传递给 prepare 方法。这包括你的训练和评估 DataLoaders、模型和优化器
>>> train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
... train_dataloader, eval_dataloader, model, optimizer
... )反向传播
最后一个添加是将训练循环中的典型 loss.backward() 替换为 🤗 Accelerate 的 backward 方法
>>> for epoch in range(num_epochs):
... for batch in train_dataloader:
... outputs = model(**batch)
... loss = outputs.loss
... accelerator.backward(loss)
... optimizer.step()
... lr_scheduler.step()
... optimizer.zero_grad()
... progress_bar.update(1)正如你在以下代码中看到的，你只需要在训练循环中添加四行额外的代码即可启用分布式训练！
+ from accelerate import Accelerator
from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler
+ accelerator = Accelerator()
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)
- device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
- model.to(device)
+ train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
+ train_dataloader, eval_dataloader, model, optimizer
+ )
num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
"linear",
optimizer=optimizer,
num_warmup_steps=0,
num_training_steps=num_training_steps
)
progress_bar = tqdm(range(num_training_steps))
model.train()
for epoch in range(num_epochs):
for batch in train_dataloader:
- batch = {k: v.to(device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
- loss.backward()
+ accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)训练
添加了相关的代码行后，在脚本或像 Colaboratory 这样的笔记本中启动你的训练。
使用脚本训练
如果你正在从脚本运行你的训练，请运行以下命令来创建和保存配置文件
accelerate config
然后使用以下命令启动你的训练：
accelerate launch train.py
使用笔记本训练
如果你计划使用 Colaboratory 的 TPU， 🤗 Accelerate 也可以在笔记本中运行。将负责训练的所有代码包装在一个函数中，并将其传递给 notebook_launcher
>>> from accelerate import notebook_launcher
>>> notebook_launcher(training_function)有关 🤗 Accelerate 及其丰富功能的更多信息，请参阅 文档。
< > 在 GitHub 上更新