Распределенное обучение с помощью TorchDistributor
В этой статье описывается, как выполнять распределенное обучение для моделей машинного обучения PyTorch с помощью TorchDistributor.
TorchDistributor — это модуль с открытым исходным кодом в PySpark, который помогает пользователям выполнять распределенное обучение с помощью PyTorch в кластерах Spark, поэтому он позволяет запускать задания обучения PyTorch в качестве заданий Spark. Она инициализирует среду и каналы связи между работниками и использует команду torch.distributed.run CLI для выполнения распределенного обучения между рабочими узлами.
API TorchDistributor поддерживает методы, показанные в следующей таблице.
| Метод и сигнатура | Description |
|---|---|
init(self, num_processes, local_mode, use_gpu) |
Создайте экземпляр TorchDistributor. |
run(self, main, *args) |
Выполняет распределенное обучение, вызывая main(**kwargs) , если основная функция является функцией и выполняет команду torchrun main *args CLI, если основной — путь к файлу. |
Требования
- Spark 3.4
- Databricks Runtime 13.0 ML или более поздней версии
Рабочий процесс разработки для записных книжек
Если процесс создания и обучения модели выполняется полностью из записной книжки на локальном компьютере или записной книжке Databricks, необходимо внести незначительные изменения, чтобы подготовить код к распределенному обучению.
Подготовка кода одного узла: подготовка и проверка кода одного узла с помощью PyTorch, PyTorch Lightning или других платформ, основанных на PyTorch/PyTorch Lightning, например API HuggingFace Trainer.
Подготовьте код для стандартного распределенного обучения: необходимо преобразовать однопроцессное обучение в распределенное обучение. Используйте этот распределенный код, охватываемый одной функцией обучения, которую можно использовать с
TorchDistributorэтим кодом.Перемещение импорта в функцию обучения: добавьте необходимые импорты, например
import torch, в функцию обучения. Это позволяет избежать распространенных ошибок выбора. Кроме того,device_idэти модели и данные должны быть привязаны к следующим образом:device_id = int(os.environ["LOCAL_RANK"])Запустите распределенное обучение: создайте экземпляр
TorchDistributorс нужными параметрами и вызов.run(*args)для запуска обучения.
Ниже приведен пример кода обучения:
from pyspark.ml.torch.distributor import TorchDistributor
def train(learning_rate, use_gpu):
import torch
import torch.distributed as dist
import torch.nn.parallel.DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler, DataLoader
backend = "nccl" if use_gpu else "gloo"
dist.init_process_group(backend)
device = int(os.environ["LOCAL_RANK"]) if use_gpu else "cpu"
model = DDP(createModel(), **kwargs)
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler)
output = train(model, loader, learning_rate)
dist.cleanup()
return output
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train, 1e-3, True)
Миграция обучения из внешних репозиториев
Если у вас есть существующая распределенная процедура обучения, хранящейся во внешнем репозитории, можно легко перейти в Azure Databricks, выполнив следующие действия.
- Импорт репозитория: импорт внешнего репозитория в виде папки Databricks Git.
- Создайте новую записную книжку инициализировать новую записную книжку Azure Databricks в репозитории.
- Запустите распределенное обучение в ячейке записной книжки, как показано
TorchDistributorниже:
from pyspark.ml.torch.distributor import TorchDistributor
train_file = "/path/to/train.py"
args = ["--learning_rate=0.001", "--batch_size=16"]
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train_file, *args)
Устранение неполадок
Распространенная ошибка рабочего процесса записной книжки заключается в том, что при выполнении распределенного обучения не удается найти или выбрать объекты. Это может произойти, когда инструкции импорта библиотеки не распределяются другим исполнителям.
Чтобы избежать этой проблемы, включите все инструкции импорта (например, import torch) в верхней части функции обучения, которая вызывается с TorchDistributor(...).run(<func>) и внутри других пользовательских функций, вызываемых в методе обучения.
Сбой NCCL: ncclInternalError: Internal check failed.
При возникновении этой ошибки во время обучения с несколькими узлами обычно возникает проблема с сетевым взаимодействием между gpu. Эта проблема возникает, когда NCCL (библиотека коллективных коммуникаций NVIDIA) не может использовать определенные сетевые интерфейсы для взаимодействия с GPU.
Чтобы устранить эту ошибку, добавьте следующий фрагмент кода обучения для использования основного сетевого интерфейса.
import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
Примеры записных книжек
В следующих примерах записной книжки показано, как выполнять распределенное обучение с помощью PyTorch.