Надоело настраивать кластер? Используйте 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.

У каждого актора есть копия 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