【PyTorch】用 Horovod 進行分散式訓練

| , | |

繼上次用 Horovod 做為 TensorFlow 分散式深度學習框架後，我這次如願改寫 PyTorch 了，雖然是另外一個案子…。

Horovod 使用

一樣直接拿著文件直接上了，大致可以分成 5 個步驟，這跟 TensorFlow Keras 倒是滿像的：

初始化

1
2
3
4
import horovod.torch as hvd

# Initialize Horovod
hvd.init()

為每個節點設置一個 GPU

1
torch.cuda.set_device(hvd.local_rank())

但，我記得官方不推薦用 set_device？

torch.cuda.set_device(device: Union[torch.device, str, int]) → None
Sets the current device.

Usage of this function is discouraged in favor of device. In most cases it’s better to use CUDA_VISIBLE_DEVICES environmental variable.


另外在看 Horovod 所提供的兩隻範例程式碼（resnet50 與 mnist）的時候，發現在 set_device 之後後都會接著：

1
torch.cuda.manual_seed(args.seed)

推測是因為每個節點在起網路的時候，初始值不盡相同，因此為了使初始值相同所以啟用了該值。


不過, 我想了一下，我只有在做分散式訓練的會希望所有網路擁有相同的初始值，所以我改了下程式，多點判斷式：

  1. 進行分散式訓練的時候。
  2. 由 master 統一設定。
1
2
3
4
5
6
7
8
if hvd.rank() == 0:
   seed = random.randint(1, 100)
   torch.manual_seed(seed)
   torch.cuda.manual_seed(seed)
   torch.cuda.manual_seed_all(seed)

# Horovod: pin GPU to local rank.
torch.cuda.set_device(hvd.local_rank())

依照分散的個數縮放 lr

文件範例這部分好像沒做，但上方的動作分解有。所以我按照之前做 TensorFlow Keras 的經驗直接用 lr 乘上節點個數。

1
2
3
4
5
6
# Build model and dataset
dataset = ...
model = ...

lr = 0.001
opt = optim.Adam(model.parameters(), lr=lr*hvd.size())


後來在 resnet50 中有看到它們實做的學習率調整，看起來是做 warmup ，不過這部分我沒跟著調整了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Horovod: using `lr = base_lr * hvd.size()` from the very beginning leads to worse final
# accuracy. Scale the learning rate `lr = base_lr` ---> `lr = base_lr * hvd.size()` during
# the first five epochs. See https://arxiv.org/abs/1706.02677 for details.
# After the warmup reduce learning rate by 10 on the 30th, 60th and 80th epochs.
def adjust_learning_rate(epoch, batch_idx):
   if epoch < args.warmup_epochs:
      epoch += float(batch_idx + 1) / len(train_loader)
      lr_adj = 1. / hvd.size() * (epoch * (hvd.size() - 1) / args.warmup_epochs + 1)
   elif epoch < 30:
      lr_adj = 1.
   elif epoch < 60:
      lr_adj = 1e-1
   elif epoch < 80:
      lr_adj = 1e-2
   else:
      lr_adj = 1e-3
   for param_group in optimizer.param_groups:
      param_group['lr'] = args.base_lr * hvd.size() * args.batches_per_allreduce * lr_adj

將 optimizer 包裝到 hvd.DistributedOptimizer 中

1
2
# Horovod: add Horovod DistributedOptimizer.
opt = hvd.DistributedOptimizer(opt, named_parameters=model.named_parameters())

在這邊我遇到的問題是，我有是使用 scheduler 去 reduced learning rate 的部分。我有找到一條 issue，是在討論這個：

雖然它的狀態是 open，但是程式碼的部分有合併回去了，所以我參考他們測試用的程式碼，把 master 目前的 learning rate 取出後傳到其他節點，其他節點接收到後設定到自己的 optimizer 中。

1
2
3
4
5
6
7
8
9
10
optimizer = ...
scheduler = ReduceLROnPlateau(optimizer, factor=lr_reduced_factor,
patience=5, mode='min')
for epoch in range(10):
   train(...)
   val_loss = validate(...)
   scheduler.step(val_loss)
   
   lr = multi_gpu_opt.make_broadcast_object(optimizer.param_groups[0]['lr'])
   optimizer.param_groups[0]['lr'] = lr

5. Broadcast

最後廣播給其他節點。

1
2
3
# Horovod: broadcast parameters & optimizer state.
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)

資料集切分

最後是資料集切分，這邊我直接仿照上次的作法採用分 epochs 數的作法。

參考資料

  1. Horovod with PyTorch 。檢自 Horovod documentation (2020-09-14)。
  2. pytorch_imagenet_resnet50.py 。檢自 horovod/horovod ｜ github (2020-09-14)。
  3. pytorch_mnist.py。檢自 horovod/horovod ｜ github (2020-09-14)。
  4. lbin (2018-09-06)。Horovod, 分布式进阶。檢自 知乎 (2020-09-14)。
  5. 王晗 (2018-08-04)。torch.manual_seed(1)是干嘛用的? - 王晗的回答。檢自 知乎 (2020-09-14)。
  6. tgaddair (2019-08-07)。Add support for broadcasting learning rate scheduler in PyTorch · Issue #1281。檢自 horovod/horovod ｜ github (2020-09-14)。
  7. Added PyTorch support for restoring optimizer state on model load and broadcast by tgaddair · Pull Request #371。檢自 horovod/horovod ｜ github (2020-09-14)。
  8. test_torch.py。檢自 horovod/horovod ｜ github (2020-09-14)。

更新紀錄

最後更新日期：2020-12-21
  • 2020-12-21 發布
  • 2020-09-14 完稿
  • 2020-09-10 起稿