简介
官方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.py
python -m torchrun my_script.py
如果你不想运行,你也可以直接将你想要的参数作为参数传递给。torchrun
accelerate 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。