Масштабирование глубокого обучения с помощью Horovod и Kubernetes29.05.2024 16:46

a1b31e47c60e1a99bfea4e6b3518248e.jpeg

Привет, Хабр!

Horovod — это фреймворк для распределенного глубокого обучения, изначально разработанный в Uber. Он позволяет масштабировать обучение моделей на сотни и тысячи GPU, сокращая время тренировки с недель до часов. Horovod поддерживает такие фреймворки, как TensorFlow, Keras, PyTorch и Apache MXNet, и легко интегрируется с существующими кодовыми базами, требуя минимум изменений.

В статье как раз и пойдет речь о том, как масштабировать модельки с помощью Horovod и Kubernetes.

Установим и настроим все необходимое

Docker необходим для контейнеризации приложений и их переносимости.

  1. Установка Docker на Ubuntu: 

    sudo apt-get update
sudo apt-get install -y apt-transport-https ca-certificates curl software-properties-common
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
sudo add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"
sudo apt-get update
sudo apt-get install -y docker-ce
sudo usermod -aG docker ${USER}

Kubernetes используется для оркестрации контейнеров.

  1. Установка Minikube (локал версия Kubernetes) на Ubuntu: 

    curl -Lo minikube https://storage.googleapis.com/minikube/releases/latest/minikube-linux-amd64
chmod +x minikube
sudo mv minikube /usr/local/bin/
minikube start --driver=docker

  2. Установка kubectl: 

    curl -LO "https://storage.googleapis.com/kubernetes-release/release/$(curl -s https://storage.googleapis.com/kubernetes-release/release/stable.txt)/bin/linux/amd64/kubectl"
chmod +x kubectl
sudo mv kubectl /usr/local/bin/

А теперь установим сам Horovod:

  1. Установка Horovod с поддержкой GPU: 

    FROM tensorflow/tensorflow:2.16.1-gpu

RUN apt-get update && apt-get install -y \
    libgl1-mesa-glx \
    libglib2.0-0

RUN pip install --upgrade pip
RUN pip install horovod

COPY . /app
WORKDIR /app

CMD ["python", "train.py"]

Для развертывания распределенного обучения с Horovod на Kubernetes нужно создать и настроить соответствующие манифесты Kubernetes.

  1. Создадим Docker-образ: 

    sudo apt-get update
sudo apt-get install -y libopenmpi-dev openmpi-bin

  2. Построение и пуш Docker-образа: 

    sudo apt-get install -y build-essential devscripts debhelper fakeroot
sudo apt-get install -y libnccl2 libnccl-dev

  3. Создание манифеста Kubernetes для MPIJob: 

    export HOROVOD_GPU_ALLREDUCE=NCCL
export HOROVOD_GPU_BROADCAST=NCCL

  4. Запуск MPIJob на Kubernetes: 

    import horovod.tensorflow as hvd
import tensorflow as tf

hvd.init()

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

# код построения модели и обучения

Для адекватной работы Horovod необходимо плюсом настроить MPI и NCCL.

  1. Установка MPI: 

    sudo apt-get update
sudo apt-get install -y libopenmpi-dev openmpi-bin

  2. Установка NCCL: 

    sudo apt-get install -y build-essential devscripts debhelper fakeroot
sudo apt-get install -y libnccl2 libnccl-dev

  3. Конфигурация NCCL:
    Убеждаемся, что среда готова к использованию NCCL: 

    export HOROVOD_GPU_ALLREDUCE=NCCL
export HOROVOD_GPU_BROADCAST=NCCL

Основные возможности horovod

Horovod использует концепции, основанные на MPI. Основные шаги инициализации:

  • hvd.init (): вызов инициализирует Horovod и подготавливает его к работе. Внутри hvd.init() Horovod вызывает MPI_Init для инициализации среды MPI.

  • hvd.size (): возвращает количество процессов (или воркеров), участвующих в распределенном обучении.

  • hvd.rank (): возвращает текущий ранг процесса. Это уникальный идентификатор для каждого воркера.

  • hvd.local_rank (): возвращает ранг процесса на текущей машине.

Horovod использует несколько коллективных операций для синхронизации данных между процессами:

  • hvd.broadcast (): синхронизирует начальное состояние переменных между всеми процессами, передавая данные из одного процесса (обычно с рангом 0) всем остальным.

  • hvd.allreduce (): выполняет операцию сокращения (например, суммирование или нахождение среднего) и распространяет результат обратно всем процессам. Это ключевая операция для усреднения градиентов при распределенном обучении.

  • hvd.allgather (): собирает данные от всех процессов и распространяет объединенный результат каждому процессу. Используется для сбора информации, такой как метрики или веса модели.

Horovod тесно интегрируется с NCCL для того, чтобы коммуникации между GPU были особо высокопроизводительными.

  • HOROVOD_GPU_ALLREDUCE : переменная окружения, указывающая на использование NCCL для операций allreduce.

  • HOROVOD_GPU_BROADCAST: переменная окружения, указывающая на использование NCCL для операций broadcast.

Horovod поддерживает асинхронные коммуникации и конвейеры для минимизации времени простоя GPU. Это достигается через использование неблокирующих вызовов MPI и стратегий оверлапа вычислений и передачи данных: 

# пример использования асинхронной передачи градиентов
hvd.allreduce_async(tensor, name)

В Horovod также есть механизмы для обнаружения сбоев и повторной инициализации процессов. При сбое одного из процессов Horovod может автоматом перезапустить его, сохранив текущее состояние обучения: 

try:
    hvd.init()
except Exception as e:
    hvd.shutdown()
    hvd.init()

Примеры применения

Horovod можно легко интегрировать в код TensorFlow для организации распределенного обучения. Простой пример обучения модельки на нескольких GPU: 

import tensorflow as tf
import horovod.tensorflow as hvd

# инициализация Horovod
hvd.init()

# установка видимых устройств GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

# построение модели
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

optimizer = tf.keras.optimizers.Adam(0.001 * hvd.size())

# обертывание оптимизатора
optimizer = hvd.DistributedOptimizer(optimizer)

model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# загрузка данных
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

# трансформация данных
train_images = train_images / 255.0
test_images = test_images / 255.0

# Callback для синхронизации начальных состояний
callbacks = [hvd.callbacks.BroadcastGlobalVariablesCallback(0)]

# обучение модели
model.fit(train_images, train_labels, batch_size=128, epochs=5, callbacks=callbacks, verbose=1 if hvd.rank() == 0 else 0)

Horovod также поддерживает PyTorch: 

import torch
import horovod.torch as hvd
from torchvision import datasets, transforms
from torch import nn, optim

# Инициализация Horovod
hvd.init()

# Настройка устройств
torch.cuda.set_device(hvd.local_rank())
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# загрузка данных
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transforms.ToTensor())
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=128, sampler=train_sampler)

# определение модели
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28 * 28, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
).to(device)

# определение оптимизатора
optimizer = optim.Adam(model.parameters(), lr=0.001 * hvd.size())

# обертывание оптимизатора
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())

# синхронизация начальных состояний
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)

# обучение модели
model.train()
for epoch in range(5):
    train_sampler.set_epoch(epoch)
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = nn.CrossEntropyLoss()(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0 and hvd.rank() == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}')

Horovod можно использовать с Keras: 

import tensorflow as tf
import horovod.tensorflow.keras as hvd

# инициализация Horovod
hvd.init()

# установка видимых устройств GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

# построение модели
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

optimizer = tf.keras.optimizers.Adam(0.001 * hvd.size())

# обертывание оптимизатора
optimizer = hvd.DistributedOptimizer(optimizer)

model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# загрузка данных
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

# трансформация данных
train_images = train_images / 255.0
test_images = test_images / 255.0

# Callback для синхронизации начальных состояний
callbacks = [hvd.callbacks.BroadcastGlobalVariablesCallback(0)]

# обучение модели
model.fit(train_images, train_labels, batch_size=128, epochs=5, callbacks=callbacks, verbose=1 if hvd.rank() == 0 else 0)

Horovod также можно использовать с Apache MXNet: 

import mxnet as mx
import horovod.mxnet as hvd

# инициализация Horovod
hvd.init()

# настройка устройств
context = mx.gpu(hvd.local_rank()) if mx.context.num_gpus() > 0 else mx.cpu()

# загрузка данных
train_data = mx.io.MNISTIter(
    image="train-images-idx3-ubyte",
    label="train-labels-idx1-ubyte",
    input_shape=(784,),
    batch_size=64,
    shuffle=True,
    flat=True
)

# определение модели
data = mx.sym.var('data')
fc1 = mx.sym.FullyConnected(data=data, num_hidden=128)
act1 = mx.sym.Activation(data=fc1, act_type="relu")
fc2 = mx.sym.FullyConnected(data=act1, num_hidden=10)
softmax = mx.sym.SoftmaxOutput(data=fc2, name='softmax')

# обертывание оптимизатора
optimizer_params = {'learning_rate': 0.01 * hvd.size()}
optimizer = mx.optimizer.create('adam', **optimizer_params)
optimizer = hvd.DistributedOptimizer(optimizer)

# настройка модели
model = mx.mod.Module(symbol=softmax, context=context)
model.fit(train_data, optimizer=optimizer, num_epoch=5)

Больше про фреймворки и инструменты Data Science и Machine Learning эксперты OTUS рассказывают в рамках практических онлайн-курсов. С полным каталогом курсов можно ознакомиться по ссылке.

© Habrahabr.ru