Распределённое обучение в ноутбуках

Important

Эта функция доступна в бета-версии.

@distributed Декоратор из Databricks AI Runtime SDK (databricks-sdk-air) — самый удобный способ запуска распределённого обучения из ноутбука Databricks. Оформите функцию обучения, вызовите её, и AI Runtime запускает её на всех GPU на узле, к которому подключён ваш ноутбук. Один и тот же код масштабируется от одного GPU до нескольких GPU без необходимости разворачивать кластер и настраивать средство распределённого запуска.

Ключевые моменты на этой странице:

• @distributed Декоратор запускает обучающую функцию для каждого GPU на вашем узле из блокнота.
• Он поддерживает PyTorch DDP, FSDP и DeepSpeed, а также переносит код с одной GPU в мульти-GPU с минимальными изменениями.
• Подключите ноутбук к ускорителю 8xH100 для полноценного обучения на нескольких GPU. По умолчанию @distributed использует все GPU на узле.

Замечание

На этой странице рассматривается распределенное обучение моделей с ноутбуков Databricks с помощью SDK Databricks AI Runtime. Чтобы отправить задачи распределённого обучения с вашей локальной машины, используйте команды Databricks CLI для AI Runtime. См. Использование CLI Databricks с AI Runtime.

Quickstart

Примеры на этой странице используют Databricks AI Runtime SDK, который предустановлен в среде GPU версии 6 и выше. В версии 5 и ниже используйте пакет serverless_gpu . См. Миграция от serverless_gpu. Чтобы подключить ноутбук к ускорителю, см. Шаг 1: Прикрепить ноутбук к GPU.

Добавьте декоратор к обучающей функции с помощью @distributed, затем вызовите её с помощью .distributed():

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

Для распределенного обучения требуется акселератор 8xH100, который подготавливает один узел с 8 GPU. По умолчанию @distributed использует все восемь видеокарт. Настройте num_accelerators, чтобы использовать меньше. accelerator_type параметр также необязателен и автоматически определяется по акселератору, к которому подключён ваш ноутбук.

Каждый вызов .distributed() создаёт запуск MLflow (или вложенный дочерний запуск, если уже есть активный запуск) и выводит ссылку в выводе ячейки. Полный рабочий пример см. в разделе полный пример.

Поддерживаемые платформы

@distributed API интегрируется с основными распределенными библиотеками обучения:

• PyTorch Distributed Data Parallel (DDP): стандартный параллелизм данных с несколькими GPU.
• Полностью шардированный параллелизм данных (FSDP): обучение больших моделей с эффективным использованием памяти.
• DeepSpeed: библиотека оптимизации Microsoft для обучения больших моделей.

Для реальных обучающих сценариев, использующих каждую библиотеку, см. примеры распределённого обучения.

Как @distributed работает декоратор

Когда вы вызываете декорированную функцию с помощью .distributed(), AI Runtime берёт на себя всю техническую сторону, которую в противном случае пришлось бы настраивать вручную с помощью средства запуска распределённых задач:

• Сериализация и разветвление: функция сериализируется и запускается на каждой выбранной видеокарте. Каждая видеокарта запускает копию функции с одними и теми же аргументами.
• Синхронизация среды: Среда Python и зависимости реплицируются на всех рангах, поэтому каждый процесс выполняет один и тот же код.
• Переменные среды ранга: стандартные переменные, такие как , LOCAL_RANK заполняются для каждого процесса. Прочитайте их в своей функции, чтобы разместить модель и данные на правильном устройстве.
• Сбор результатов: Возвращаемые значения собираются со всех рангов и возвращаются вызывающему.
• Отслеживание MLflow: Каждый .distributed() вызов создаёт MLflow run или вложенный дочерний запуск, если он уже активен, поэтому метрики, зарегистрированные из вашей функции, попадают в тот же запуск.
• Жизненный цикл и тайм-аут: Распределённое выполнение происходит в рамках жизненного цикла ноутбука. При остановке блокнота выполнение также прерывается. У декоратора по умолчанию тайм-аут составляет 3 часа. Укажите timeout в секундах, чтобы изменить это, или timeout=None, чтобы отключить это. Кастомные тайм-ауты требуют GPU среды v5 и выше.

API основан на стандартных библиотеках PyTorch: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) и DeepSpeed.

Перейти с serverless_gpu

Пакет serverless_gpu устарел и больше не получает новых функций. Начиная с 6-й версии среды GPU, Databricks AI Runtime SDK (databricks-sdk-air) заменяет его и предоставляет эквивалентные API с переименованными аргументами в пространстве имён databricks.air. Используемый вами пакет зависит от версии среды GPU:

• Версия 5 и ниже: включить только serverless_gpu. Продолжайте использовать serverless_gpu в этих версиях.
• Версия 6: включает оба пакета, но обновления получает только AI Runtime SDK. Версия 6 — последняя версия среды GPU, которая включает serverless_gpu.

AI Runtime SDK — это отдельный от Databricks SDK for Python пакет (databricks-sdk).

	serverless_gpu (не рекомендуется)	SDK среды выполнения ИИ (databricks.air)
Версии среды GPU	6 и ниже	6 и старше
Импорт	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Аргумент количества GPU	gpus	num_accelerators
Аргумент типа графического процессора	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Значения типа GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Модуль загрузки данных и создания контрольных точек	serverless_gpu.data	databricks.air.data. См. Загрузка данных в среде выполнения ИИ.

Чтобы перенести блокнот:

1. Подключите ноутбук к среде GPU версии 6 или выше, которая включает AI Runtime SDK.
2. Измените импорт из serverless_gpu в databricks.air.
3. Переименуйте gpus в num_accelerators и gpu_type в accelerator_type, используйте новые значения типа ускорителя. Старые имена аргументов и значения типа GPU всё ещё работают в databricks.air, но выдают FutureWarning. Поддержка прекратится в GPU Environment версии 7.

Tip

Попросите Genie Code сделать это за вас:

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

Для рабочих нагрузок Databricks CLI, использующих databricks.air.data, добавьте databricks-sdk-air[data] в environment.dependencies рабочей нагрузки.

Поступает от TorchDistributor

Если вы запускаете распределённое обучение с помощью TorchDistributor на кластерах Spark и ваша нагрузка помещается на одном узле, @distributed API является рекомендуемой заменой в среде AI Runtime. Код вашей модели DDP, FSDP или DeepSpeed не меняется.

Feature	@distributed API-интерфейс	Распространитель факелов
Infrastructure	Полностью бессерверное решение, без управления кластерами	Требуется кластер Spark со службами, работающими на GPU.
Setup	Один декоратор, минимальная конфигурация	Требуется настройка кластера Spark и TorchDistributor
Поддержка фреймворка	PyTorch DDP, FSDP, DeepSpeed	В первую очередь PyTorch DDP
Загрузка данных	Внутри декоратора используются тома Unity Catalog (UCVolumeDataset для потоковой передачи файлов)	С помощью Spark или файловой системы

Для миграции одноузловой рабочей нагрузки:

• Замените вызов TorchDistributor(...).run(train_fn, ...) на декоратор @distributed для train_fn, затем запустите с помощью train_fn.distributed(...).
• Удалите кластер Spark и конфигурацию GPU worker. Вместо этого подключите ноутбук к ускорителю 8xH100. @distributed по умолчанию использует все восемь видеокарт.
• Переместите загрузку данных внутрь декорированной функции. См. загрузку данных.
• Сохраняйте существующий код модели DDP, FSDP или DeepSpeed. Декоратор поддерживает все три варианта.

Декоратор @distributed находится в бета-версии и работает на одном узле (см. Ограничения), поэтому он не заменяет все рабочие нагрузки TorchDistributor. Оставьте рабочие нагрузки, требующие интеграции со Spark, на TorchDistributor. Для запуска распределённого обучения с локального компьютера или между несколькими узлами используйте команды Databricks CLI для AI Runtime. См. Использование CLI Databricks с AI Runtime.

Полный пример

Следующий пример обучает многослойную модель перцептрона (MLP) на 8 графических процессорах H100 из ноутбука.

1. Настройте модель и определите служебные функции.

   
   # Define the model
   import os
   import torch
   import torch.distributed as dist
   import torch.nn as nn
   
   def setup():
       torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
       dist.init_process_group("nccl")
   
   def cleanup():
       dist.destroy_process_group()
   
   class SimpleMLP(nn.Module):
       def __init__(self, input_dim=10, hidden_dim=64, output_dim=1):
           super().__init__()
           self.net = nn.Sequential(
               nn.Linear(input_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, output_dim)
           )
   
       def forward(self, x):
           return self.net(x)
   

2. Импортируйте distributed декоратор из AI Runtime SDK.

   from databricks.air import distributed
   

3. Оберните код обучения модели в функцию и декорируйте функцию декоратором @distributed. Украшенная функция является входной точкой распределённого выполнения, поэтому задайте всю обучающую логику, загрузку данных и инициализацию модели внутри неё.

   @distributed()
   def run_train(num_epochs: int, batch_size: int) -> None:
       import mlflow
       import torch.optim as optim
       from torch.nn.parallel import DistributedDataParallel as DDP
       from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
   
       # 1. Set up multi-GPU environment
       setup()
       device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
   
       # 2. Apply the Torch distributed data parallel (DDP) library for data-parallel training.
       model = SimpleMLP().to(device)
       model = DDP(model, device_ids=[device])
   
       # 3. Create and load dataset.
       x = torch.randn(5000, 10)
       y = torch.randn(5000, 1)
   
       dataset = TensorDataset(x, y)
       sampler = DistributedSampler(dataset)
       dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
   
       # 4. Define the training loop.
       optimizer = optim.Adam(model.parameters(), lr=0.001)
       loss_fn = nn.MSELoss()
   
       for epoch in range(num_epochs):
           sampler.set_epoch(epoch)
           model.train()
           total_loss = 0.0
           for step, (xb, yb) in enumerate(dataloader):
               xb, yb = xb.to(device), yb.to(device)
               optimizer.zero_grad()
               loss = loss_fn(model(xb), yb)
               # Log loss to MLflow metric
               mlflow.log_metric("loss", loss.item(), step=step)
   
               loss.backward()
               optimizer.step()
               total_loss += loss.item() * xb.size(0)
   
           mlflow.log_metric("total_loss", total_loss)
           print(f"Total loss for epoch {epoch}: {total_loss}")
   
       cleanup()
   

4. Запускайте распределённое обучение, вызывая распределённую функцию с пользовательскими аргументами.

   run_train.distributed(num_epochs=3, batch_size=1)
   

5. При выполнении в выходных данных ячейки блокнота появляется ссылка на запуск MLflow. Щелкните ссылку запуска MLflow или найдите ее на панели "Эксперимент ", чтобы просмотреть результаты выполнения. Дополнительные сведения о настройке имен экспериментов, отслеживания метрик и возобновлении выполнения см. в разделе "Отслеживание экспериментов" и "Наблюдаемость".

Загрузка данных

Поместите код загрузки данных внутри функции @distributed. Набор данных может превышать максимальный размер, допустимый для pickle, поэтому создание или загрузка его внутри декоратора позволяет избежать ошибок сериализации:

from databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Для файловых данных, хранящихся в томах Unity Catalog, используйте UCVolumeDataset из databricks.air.data, который передаёт файлы в потоковом режиме с локальным кэшированием и автоматически распределяет их между рангами и рабочими процессами. Для контрольных точек распределенного обучения в том, используйте UCVolumeWriter и UCVolumeReader. См. Загрузка данных в AI Runtime и Checkpoint с распределённой контрольной точкой (DCP).

Limitations

• @distributed поддерживает только один узел: он работает между GPU на узле, к которому подключён ваш блокнот. Для обучения на нескольких узлах используйте CLI AI Runtime Databricks. См. Использование CLI Databricks с AI Runtime.
• Если вы установите accelerator_type, он должен совпадать с акселератором, к которому подключён ваш ноутбук (например, "GPU_8xH100"). Несоответствие приводит к сбою рабочей нагрузки. Параметр необязательный и автоматически определяется при его отсутствии.
• AI Runtime поддерживает GPU Environment v4 и выше. Настраиваемые тайм-ауты (параметр timeout) требуют среды GPU версии v5 или выше.
• По умолчанию тайм-аут декоратора составляет 3 часа. Укажите timeout в секундах, чтобы изменить это, или timeout=None, чтобы отключить это.
• Выполнение происходит в рамках жизненного цикла ноутбука. При остановке блокнота выполнение также прерывается.

Дополнительные ресурсы

• Для декоратора @distributed, AcceleratorType и API Ray см. справочник по API Databricks AI Runtime SDK.
• Для паттернов, которые делают ваш учебный процесс более эффективным и устойчивым, ознакомьтесь с руководством по эффективности и устойчивости.
• Для сквозных сценариев обучения см. примеры распределённого обучения.