Пример обучения модели с использованием библиотеки Horovod

Мы предлагаем использовать библиотеку Horovod в связи с простотой адаптации нераспределенного кода в код, исполняемый на кластере, и высокой эффективностью распараллеливания.

Horovod — это библиотека, разработанная компанией Uber, которая позволяет преобразовывать код для выполнения на одной машине в код, исполняемый на кластере, путем добавления нескольких строк кода.

Как обучать модели с помощью Horovod?

Далее будут рассмотрены примеры использования Horovod с библиотеками TensorFlow и PyTorch.

TensorFlow и PyTorch

  1. Выполните команду hvd.init() в самом начале скрипта.

  2. Прикрепите GPU к процессу MPI. При типичной настройке одного графического процессора на процесс это может быть выполнено установкой на hvd.local_rank(). В этом случае первому процессу на сервере будет выделен первый графический процессор, второму процессу — второй графический процессор и так далее.

    Tensorflow:

    config = tf.ConfigProto()
    config.gpu_options.visible_device_list = str(hvd.local_rank())
    

    PyTorch:

    torch.cuda.set_device(hvd.local_rank())
    
  3. Адаптируйте темп обучения (learning rate) под изменившийся размер батча. Простейший способ — умножить старое значение на число процессов MPI. Однако существуют и более сложные стратегии.

    Tensorflow:

    opt = tf.train.AdagradOptimizer(0.01 * hvd.size))
    

    PyTorch:

    optimizer = optim.SGD(model.parameters(), lr=args.lr * hvd.size(), momentum=args.momentum)
    

    Оберните свой оптимизатор классом hvd. DistributedOptimizer

    Tensorflow:

    opt = hvd.DistributedOptimizer(opt)
    

    PyTorch:

    optimizer = hvd.DistributedOptimizer(optimizer,
    named_parameters=model.named_parameters())
    
  4. Добавьте hvd.BroadcastGlobalVariablesHook (0) для передачи начальных состояний переменных с процесса 0 всем другим процессам. Это необходимо для обеспечения последовательной инициализации всех процессов, когда обучение начинается со случайных весов или восстанавливается с контрольной точки.

    hvd.broadcast_global_variables(0)
    

    или

    hooks = [hvd.BroadcastGlobalVariablesHook(0)]
    

    и далее передать хук аргументом в MonitoredTrainingSession, TrainingSession, estimator или иной потребитель, который вы используете для обучения.

    pythonhvd.broadcast_parameters(model.state_dict(), root_rank=0)
    
  5. Измените логику записи контрольных точек так, чтобы она происходила только на процессе с номером 0, например:

    if hvd.rank() == 0: write smth
    

    Опционально, но важно: используйте распределенные генераторы данных, например:

    PyTorch

    # 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=batch_size, sampler=train_sampler)
    

    Если в TensorFlow вы используете градиенты напрямую, использовать команду grads = tf.gradients(loss, tvars) не получится. Вместо этого замените код на следующий:

    TensorFlow

    grads_and_vars = optimizer.compute_gradients(loss, tvars)
    grads = [grad for grad, var in grads_and_vars]
    

    Рекомендуется проверять данные с использованием только одного процесса:

    hvd.rank() == 0.
    

    Запуск и отладка скриптов через Jupyter Notebook:

    Выберите один из образов с пометкой horovod (прим. jupyter-horovod-tf15)

    Запуск с Horovod:

    mpirun -np {GPU count} python train_horovod_example.py
    

    Запуск с Pytorch.distributed:

    python -m torch.distributed.launch --nproc_per_node {GPU count} train_distributed_example.py