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

简介

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

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

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

示例如下

  import torchimport torch.nn.functional as Ffrom datasets import load_dataset
+ from accelerate import Accelerator+ accelerator = Accelerator()
- device = 'cpu'
+ device = accelerator.devicemodel = 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 torchimport torch.nn.functional as Ffrom 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()