Надоело настраивать кластер? Используйте Ray Lightning

PyTorch Lightning позволяет распараллелить Deep Learning на GPU, но настраивать и объединять процессоры в сеть сложно даже в управляемом кластере SLURM. Проблему решает пакет Ray Lightning, обзором которого делимся к старту потока курса по Data Science.

---

Напомним: когда вы настраиваете многоузловое обучение на GPU, нужно, чтобы каждый узел:

• взаимодействовал со всеми остальными;

• имел доступ к коду;

• имел корректно настроенные переменные PyTorch;

• также нужно запустить сценарий обучения на каждом узле;

И вот с какими ограничениями PyTorch Lightning вы столкнётесь:

• настройка кластера на любом облаке (AWS, Azure, GCP или Kubernetes) требует широких компетенций;

• многоузловое обучение не поддерживается Jupyter Notebook;

• автоматическое вертикальное масштабирование, чтобы снизить затраты, потребует обширной инфраструктуры и специальных инструментов.

Можно ли обойтись без глубоких знаний построения большой инфраструктуры, а также без изменений кода?

Знакомьтесь с Ray Lightning

Ray Lightning — это простой плагин PyTorch Lightning для горизонтального масштабирования обучения. Вот его преимущества:

• Его просто настраивать, не придётся изменять код.

• Легко масштабировать: пишем один код для одного GPU и меняем один параметр для перехода на больший кластер.

• Он работает с Jupyter Notebook: запустив Ray Lightning, вы получите доступ ко всему кластеру.

• При помощи Ray Cluster Launcher легко настроить многоузловой кластер на AWS/Azure/GCP.

• Для крупномасштабного распределённого поиска гиперпараметров он интегрируется с Ray Tune и алгоритмами SOTA.

• А главное — он бесплатный, его исходный код открыт!

Под капотом Ray Lightning — Ray, простая библиотека распределённых вычислений на Python.

Как работает Ray Lightning?

Ray Lightning задействует интерфейс плагина PyTorch Lightning, чтобы предложить RayPlugin, который вы можете добавить в Trainer. Он работает аналогично плагину DDPSpawn, но вместо новых процессов для обучения создаёт новые акторы Ray. Акторы — это также процессы Python, но их можно запланировать в любом месте кластера Ray, что позволяет программировать, не выходя из сценария Python.

Взаимодействие акторов Ray в многоузловом кластере

У каждого актора есть копия LightningModule, акторы автоматически устанавливают корректные переменные окружения и вместе создают группу взаимодействия PyTorch. Ray запускает стандартный DistributedDataParallel с той же производительностью, но обучение выполняется программно, а экземпляры по мере обучения автоматически масштабируются вертикально.

Управление кластером

Обычно управление кластерами без команды, работающей с платформой машинного обучения или инфраструктурой, — это мучение. Но только не с Ray, где кластер запускается специальным инструментом.

Этоn инструмент поддерживает AWS, GCP, Azure, и в нём есть оператор Kubernetes, поэтому программа на Ray запускается где угодно. Выполняемый в кластере Ray код легко мигрирует или переносится с одного облака на другое. Чтобы запустить кластер Ray на AWS, нужен конфигурационный файл YAML:

cluster_name: ml

# Cloud-provider specific configuration.
provider:
    type: aws
    region: us-west-2
    availability_zone: us-west-2a,us-west-2b

# How Ray will authenticate with newly launched nodes.
auth:
    ssh_user: ubuntu

head_node:
    InstanceType: p3.8xlarge
    ImageId: latest_dlami

    # You can provision additional disk space with a conf as follows
    BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
              VolumeSize: 100
worker_nodes:
    InstanceType: p3.2xlarge
    ImageId: latest_dlami

file_mounts: {
    "/path1/on/remote/machine": "/path1/on/local/machine",
}

# List of shell commands to run to set up nodes.
setup_commands:
    - pip install -U ray-lightning

Данные в file_mounts синхронизируются со всеми узлами кластера, поэтому обучающий сценарий помещаем сюда. Любые дополнительно устанавливаемые зависимости (например, pip install foobar) указываем в setup_commands. Они установятся на всех узлах кластера.

Заполучив файл YAML, выполняем ray up cluster.yaml, чтобы создать и запустить кластер. Выполняем ray attach cluster.yaml, чтобы подключиться по ssh к головному узлу кластера Ray.

Когда ресурсов запрашивается больше, чем доступно в текущем кластере, средство запуска кластеров автоматически добавляет новые узлы, автоматически удаляя неиспользуемые.

Теперь соберём всё

Посмотрим, как легко обучить простой классификатор MNIST в облаке.

Установка

Установите Ray Lightning:

pip install ray-lightning

Также будут установлены PyTorch Lightning и Ray.

Вернёмся к PyTorch Lightning

Подготовим код PyTorch Lightning. Создадим модель классификатора — экземпляр LightningModule. Вот пример простого классификатора MNIST из руководства по PyTorch Lightning:

import pytorch_lightning as pl
import torch
from torch.utils.data import random_split, DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms

class LightningMNISTClassifier(pl.LightningModule):
    def __init__(self, config, data_dir=None):
        super(LightningMNISTClassifier, self).__init__()

        self.data_dir = data_dir
        self.lr = config["lr"]
        layer_1, layer_2 = config["layer_1"], config["layer_2"]
        self.batch_size = config["batch_size"]

        # mnist images are (1, 28, 28) (channels, width, height)
        self.layer_1 = torch.nn.Linear(28 * 28, layer_1)
        self.layer_2 = torch.nn.Linear(layer_1, layer_2)
        self.layer_3 = torch.nn.Linear(layer_2, 10)
        self.accuracy = pl.metrics.Accuracy()

    def forward(self, x):
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = torch.relu(x)
        x = self.layer_2(x)
        x = torch.relu(x)
        x = self.layer_3(x)
        x = F.softmax(x, dim=1)
        return x

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)

    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        logits = self.forward(x)
        loss = F.nll_loss(logits, y)
        acc = self.accuracy(logits, y)
        self.log("ptl/train_loss", loss)
        self.log("ptl/train_accuracy", acc)
        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y = val_batch
        logits = self.forward(x)
        loss = F.nll_loss(logits, y)
        acc = self.accuracy(logits, y)
        return {"val_loss": loss, "val_accuracy": acc}

    def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        avg_acc = torch.stack([x["val_accuracy"] for x in outputs]).mean()
        self.log("ptl/val_loss", avg_loss)
        self.log("ptl/val_accuracy", avg_acc)

    def prepare_data(self):
        self.dataset = MNIST(
            self.data_dir,
            train=True,
            download=True,
            transform=transforms.ToTensor())

    def train_dataloader(self):
        dataset = self.dataset
        train_length = len(dataset)
        dataset_train, _ = random_split(
            dataset, [train_length - 5000, 5000],
            generator=torch.Generator().manual_seed(0))
        loader = DataLoader(
            dataset_train,
            batch_size=self.batch_size,
            num_workers=1,
            drop_last=True,
            pin_memory=True,
        )
        return loader

    def val_dataloader(self):
        dataset = self.dataset
        train_length = len(dataset)
        _, dataset_val = random_split(
            dataset, [train_length - 5000, 5000],
            generator=torch.Generator().manual_seed(0))
        loader = DataLoader(
            dataset_val,
            batch_size=self.batch_size,
            num_workers=1,
            drop_last=True,
            pin_memory=True,
        )

Создадим экземпляр модели и Trainer, обучим медель:

model = LightningMNISTClassifier(config, data_dir="./")

trainer = pl.Trainer(max_epochs=10)
trainer.fit(model)

Этого достаточно для однопоточного выполнения. Приступим к обучению на ноутбуке. Распараллелим процесс на большой кластер, используя GPU и Ray Lightning.

Параллельное выполнение на ноутбуке

Распараллелим обучение между ядрами ноутбука, добавив RayPlugin в Trainer и отключив графические процессоры:

import ray
from ray_lightning import RayPlugin

class LightningMNISTClassifier(...):
   # ... etc

# variables for Ray around parallelism and hardware
num_workers = 8
use_gpu = False

# Initialize ray.
ray.init()

model = LightningMNISTClassifier(config, data_dir)

trainer = pl.Trainer(
    max_epochs=10, 
    plugins=[RayPlugin(num_workers=num_workers, use_gpu=use_gpu)])
trainer.fit(model)

Снова запускаем сценарий, небольшими изменениями распределив обучение на 8 воркеров, то есть процессов.

Обучение на множестве узлов кластере Ray

Чтобы распараллелить обучение на кластере с несколькими графическими процессорами и узлами, задействуем средство запуска кластеров Ray с RayPlugin. Запускаем кластер:

ray up cluster.yaml

В file_mounts обязательно вставить обучающий сценарий, а также прописать все зависимости pip в разделе setup_commands. Полное пошаговое руководство со всевозможными конфигурациями файла YAML, смотрите здесь.

Запустив кластер, по SSH подключаемся к головному узлу:

ray attach cluster.yaml

Добавленный в file_mounts из cluster.yaml обучающий сценарий синхронизируется с этим головным узлом. В том же коде меняем две строки:

• ray.init()-> ray.init("auto") чтобы Ray подключился к кластеру, а не запустил локальный экземпляр.

• use_gpu устанавливаем в True, а в num_workers указываем общее число процессов/GPU.

И выполняем скрипт:

python train.py

Не конфигурируя кластер и почти не изменяя код, мы запустили задачу обучения с распределёнными данными, несколькими узлами и GPU.

Заключение

Мы подготовили всё для многоузлового обучения на GPU, но Ray Lightning не ограничивается этим:

• Если стандартный PyTorch DDP вам не по душе, попробуйте RayHorovodPlugin с Horovod вместо DDP для базового протокола распределённого обучения, а также RayShardedPlugin — (эффективная в смысле модельпараллельного обучения с Fairscale).

• Ray Lightning также интегрируется с Ray Tune, позволяя проводить эксперименты по настройке распределённых гиперпараметров с каждым прогоном распараллеленного обучения. Подробности в полном руководстве по Ray+PyTorch Lightning E2E.

• Для обучения в облаке используйте Ray Client. Подробности здесь.

Попробовать Ray Lightning в деле вы сможете на нашем курсе по Data Science, в конце которого ждёт специализация в Machine Learning. Также вы можете перейти на страницы других курсов, чтобы узнать, как мы готовим специалистов в других направлениях:

Data Science и Machine Learning

• Профессия Data Scientist

• Профессия Data Analyst

• Курс «Математика для Data Science»

• Курс «Математика и Machine Learning для Data Science»

• Курс по Data Engineering

• Курс «Machine Learning и Deep Learning»

• Курс по Machine Learning

Python, веб-разработка

• Профессия Fullstack-разработчик на Python

• Курс «Python для веб-разработки»

• Профессия Frontend-разработчик

• Профессия Веб-разработчик

Мобильная разработка

• Профессия iOS-разработчик

• Профессия Android-разработчик

Java и C#

• Профессия Java-разработчик

• Профессия QA-инженер на JAVA

• Профессия C#-разработчик

• Профессия Разработчик игр на Unity

От основ — в глубину

• Курс «Алгоритмы и структуры данных»

• Профессия C++ разработчик

• Профессия Этичный хакер

А также:

• Курс по DevOps