Accelerate
Accelerate 是一个库,通过添加四行代码,就可以让相同的 PyTorch 代码在任何分布式配置中运行!简而言之,它可以让大规模训练和推理变得简单、高效且适应性强。
+ from accelerate import Accelerator
+ accelerator = Accelerator()
+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+ model, optimizer, training_dataloader, scheduler
+ )
for batch in training_dataloader:
optimizer.zero_grad()
inputs, targets = batch
inputs = inputs.to(device)
targets = targets.to(device)
outputs = model(inputs)
loss = loss_function(outputs, targets)
+ accelerator.backward(loss)
optimizer.step()
scheduler.step()
Accelerate 基于 torch_xla
和 torch.distributed
,它负责繁重的任务,因此您无需编写任何自定义代码来适应这些平台。将现有的代码库转换为利用 DeepSpeed,执行 全分片数据并行,并自动支持混合精度训练!
要更好地了解这个过程,请务必查看 教程!
然后,您可以通过 Accelerate 的 CLI 界面在任何系统上启动此代码。
accelerate launch {my_script.py}
< > 更新 在 GitHub 上