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

news2024/9/24 17:33:15

简介

官方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()