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

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

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

Примеры использования

Рассмотрим примеры использования Horovod с библиотеками 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. Адаптируйте темп обучения под изменившийся размер батча. Простейший способ — умножить старое значение на число процессов 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

  1. Выберите один из образов с библиотекой horovod в составе, например cr.msk.sbercloud.ru/aicloud-jupyter/jupyter-cuda11.5-pt1.11.0-gpu.

  2. Запустите скрипт обучения.

    Запуск с 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
    
Масштабная конференция
GoCloud 2024:
облачные грани будущего