利用Accelerate()进行pytorch的多GPU加速

简介

官方Github：https://github.com/huggingface/accelerate

Accelerate 是为喜欢编写PyTorch模型的训练循环但不愿意编写和维护使用多GPU/TPU/fp16所需的样板代码的PyTorch用户创建的。

它可以仅加速与多 GPU/TPU/fp16 相关的样板代码，并保持其余代码不变。

示例如下

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

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

          optimizer.step()

正如您在此示例中所看到的，通过向任何标准 PyTorch 训练脚本添加 5 行代码，您现在可以在任何类型的单节点或分布式节点设置（单 CPU、单 GPU、多 GPU 和 TPU）以及使用或不使用混合精度（fp8、fp16、bf16）上运行。

特别是，无需修改即可在本地计算机上运行相同的代码，以进行调试或训练环境。

Accelerate 甚至会为您处理设备放置（这需要对代码进行更多更改，但通常更安全），因此您甚至可以进一步简化训练循环：

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

- device = 'cpu'
+ accelerator = Accelerator()

- model = torch.nn.Transformer().to(device)
+ 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:
-         source = source.to(device)
-         targets = targets.to(device)

          optimizer.zero_grad()

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

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

          optimizer.step()

启动脚本：

 accelerate launch my_script.py --args_to_my_script

此 CLI 工具是可选的，您仍然可以使用或在方便时使用。python my_script.pypython -m torchrun my_script.py

如果你不想运行，你也可以直接将你想要的参数作为参数传递给。torchrunaccelerate launch accelerate config

例如，以下是在两个 GPU 上启动的方法：

accelerate launch --multi_gpu --num_processes 2 examples/nlp_example.py

更多查看官方提供的CLI文档：https://huggingface.co/docs/accelerate/package_reference/cli

如果想指定GPU，可以通过在终端运行 accelerate config 命令进行配置，根据需求会最终生成一个 default_config.yaml 文件如下：

这里面就可以设置对应的gpu_ids。