Transformers 文档

使用🤗 Accelerate进行分布式训练

使用🤗 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的TPUs,Accelerate也可以在笔记本中运行。将所有负责训练的代码包装在一个函数中,并将其传递给notebook_launcher

>>> from accelerate import notebook_launcher

>>> notebook_launcher(training_function)

有关🤗 Accelerate及其丰富功能的更多信息,请参阅文档

< > Update on GitHub