
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.

У каждого актора есть копия 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
Python, веб-разработка
Мобильная разработка
Java и C#
От основ — в глубину
А также:
 
           
 
Dilirious
Вопрос немного не по теме, но мне никак не понять как это сделать.
В браузере вводим habr.com. Пояснить схему работы DNS и показать, как проходит запрос, откуда откроется сайт и где расположены файлы сайта (имя сервера).