Horovod 与 PyTorch¶

要使用 Horovod 与 PyTorch，请对您的训练脚本进行以下修改：

运行 hvd.init()。

将每个GPU固定分配给单个进程。

在典型的每个进程一个GPU的设置中，将此设置为本地排名。服务器上的第一个进程将被分配第一个GPU，第二个进程将被分配第二个GPU，依此类推。
```
if torch.cuda.is_available():
    torch.cuda.set_device(hvd.local_rank())
```

根据工作节点数量缩放学习率。

在同步分布式训练中，有效批次大小会随工作节点数量成比例增加。提高学习率可以补偿增大的批次大小。

将优化器包装在hvd.DistributedOptimizer中。

分布式优化器将梯度计算委托给原始优化器，使用allreduce或allgather对梯度进行平均，然后应用这些平均后的梯度。

将初始变量状态从rank 0广播到所有其他进程：
```
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
```
当使用随机权重开始训练或从检查点恢复时，这对于确保所有工作节点的一致初始化是必要的。

修改你的代码，仅在工作者0上保存检查点，以防止其他工作者损坏它们。

通过使用hvd.rank() != 0来保护模型检查点代码实现这一点。

示例（另请参见完整训练示例）：

import torch
import horovod.torch as hvd

# Initialize Horovod
hvd.init()

# Pin GPU to be used to process local rank (one GPU per process)
torch.cuda.set_device(hvd.local_rank())

# Define dataset...
train_dataset = ...

# Partition dataset among workers using DistributedSampler
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

# Build model...
model = ...
model.cuda()

optimizer = optim.SGD(model.parameters())

# Add Horovod Distributed Optimizer
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())

# Broadcast parameters from rank 0 to all other processes.
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
       optimizer.zero_grad()
       output = model(data)
       loss = F.nll_loss(output, target)
       loss.backward()
       optimizer.step()
       if batch_idx % args.log_interval == 0:
           print('Train Epoch: {} [{}/{}]\tLoss: {}'.format(
               epoch, batch_idx * len(data), len(train_sampler), loss.item()))

注意

PyTorch GPU支持需要NCCL 2.2或更高版本。如果不使用RoCE或InfiniBand，也适用于NCCL 2.1.15。

PyTorch Lightning¶

Horovod 作为分布式后端在 PyTorch Lightning 中从 v0.7.4 及以上版本开始得到支持。

使用 PyTorch Lightning，通过 Horovod 进行分布式训练仅需对现有训练脚本进行单行代码修改：

# train Horovod on GPU (number of GPUs / machines provided on command-line)
trainer = pl.Trainer(accelerator='horovod', gpus=1)

# train Horovod on CPU (number of processes / machines provided on command-line)
trainer = pl.Trainer(accelerator='horovod')

在某些旧版本的pytorch_lightning中，可能需要将参数“accelerator”的名称更改为“distributed_backend”。

启动训练任务，并像平常使用Horovod时那样在命令行中指定工作进程数量：

# run training with 4 GPUs on a single machine
$ horovodrun -np 4 python train.py

# run training with 8 GPUs on two machines (4 GPUs each)
$ horovodrun -np 8 -H hostname1:4,hostname2:4 python train.py

你可以在pytorch_lightning_mnist.py中找到使用horovod后端的pytorch lightning训练器的示例

查看 PyTorch Lightning 文档获取更多详情。

还添加了一个基于Pytorch-Lightning的spark估计器，示例位于pytorch_lightning_spark_mnist.py

Horovod 与 PyTorch¶

PyTorch Lightning¶

导航

相关主题