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 и показать, как проходит запрос, откуда откроется сайт и где расположены файлы сайта (имя сервера).