使用 🤗 Accelerate 进行分布式训练¶
随着模型变得越来越大，并行性已经成为在有限硬件上训练更大模型和加速训练速度的策略，增加了数个数量级。
在 Hugging Face中我们创建了🤗 加速库 Accelerate，以帮助用户在任何类型的分布式设置上轻松训练🤗 Transformers模型，无论是在一台机器上的多个GPU还是在多个机器上的多个GPU。
在本教程中，将了解如何自定义原生 PyTorch 训练循环，以启用在分布式环境中的训练。
设置¶
通过安装 🤗 Accelerate 开始:
pip install accelerate
然后导入并创建Accelerator对象。Accelerator 将自动检测分布式设置的类型，并初始化所有必要的训练组件。你不需要显式地设置将模型放在哪个设备上。
from accelerate import Accelerator
accelerator = Accelerator()
准备加速¶
下一步是将所有相关的训练对象传递给 prepare 方法。这包括训练和评估的DataLoader、一个模型和一个优化器:
train_dataloader用于训练阶段，加载数据并分批提供给模型进行训练。eval_dataloader用于评估阶段，加载数据并分批提供给模型进行性能评估。
accelerator.prepare 的作用： 在 Hugging Face Accelerate 库中，accelerator.prepare 方法用于将 DataLoader、模型和优化器等对象包装起来，以便在分布式训练环境中正确地管理它们。具体作用包括：
- 自动分发：将数据、模型和优化器自动分发到不同的 GPU 或 TPU 上。
- 同步：在分布式训练中，确保各个进程之间的数据同步。
- 优化：提供性能优化，如自动混合精度训练。
train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
train_dataloader, eval_dataloader, model, optimizer
)
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)
# Add patch #
train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
train_dataloader, eval_dataloader, model, optimizer
)
# Add patch #
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) # Add
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
训练¶
在添加了相关代码行后，可以在脚本或 notebook（如Colaboratory）中启动训练。
用脚本训练¶
如果在脚本中运行训练，请运行以下命令以创建和保存配置文件:
accelerate config
然后使用以下命令启动训练:
accelerate launch train.py
用 notebook 训练¶
使用 Colaboratory 的 TPU 支持在 notebook 中运行 Accelerate。它将所有负责训练的代码包装在一个函数中，并将其传递给 notebook_launcher：
from accelerate import notebook_launcher
notebook_launcher(training_function)
有关 🤗 Accelerate 及其丰富功能等更多信息，可参阅文档。