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

Это важно

Эта функция доступна в бета-версии. Администраторы рабочей области могут управлять доступом к этой функции на странице "Предварительные версии ". См. статью "Управление предварительными версиями Azure Databricks".

Замечание

Эта страница посвящена распределенному обучению в ноутбуках Databricks с использованием бессерверного Python API для GPU. Чтобы отправлять задачи распределенного обучения с локального компьютера, используйте CLI AI Runtime, который доступен в публичной предварительной версии. См. интерфейс командной строки среды выполнения ИИ.

Вы можете запускать распределенные рабочие нагрузки между несколькими GPU на одном узле с помощью API Serverless GPU Python. API предоставляет простой унифицированный интерфейс, который абстрагирует сведения о подготовке GPU, настройке среды и распределении рабочей нагрузки. При минимальных изменениях кода вы можете легко перейти от обучения с одним GPU к распределенному выполнению с несколькими GPU из одной записной книжки.

Замечание

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

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

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

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

Serverless_GPU API и TorchDistributor

В следующей serverless_gpu@distributed таблице сравнивается API с TorchDistributor:

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

API serverless_gpu — это рекомендуемый подход для новых рабочих нагрузок глубокого обучения в Databricks. TorchDistributor остается доступным для рабочих нагрузок, тесно связанных с кластерами Spark.

Быстрый старт

Бессерверный API GPU для распределенного обучения предварительно установлен при подключении к бессерверному GPU в записных книжках и заданиях Databricks. Рекомендуется использовать среду GPU 4 или выше. Чтобы использовать его для распределенного distributed обучения, импортируйте и используйте декоратор для распространения функции обучения.

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

Чтобы запустить распределенное выполнение, вызовите декорированную функцию с помощью train_function.distributed(). Каждый вызов автоматически создает запуск эксперимента MLflow или вложенный дочерний запуск, если он уже активен.

Предупреждение

Если вы задаёте gpu_type в @distributed, убедитесь, что его значение соответствует типу акселератора, к которому подключён ваш ноутбук ("H100" или "A10"). Указание неправильного типа акселератора приведет к сбою рабочей нагрузки.

В приведенном ниже фрагменте кода показано базовое использование @distributed.

from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs.
# gpu_type='H100' is optional and will be auto-detected if not set.
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

run_train.distributed()

Ниже приведен полный пример, который обучает многоуровневую модель perceptron (MLP) на 8 GPU H100 из записной книжки:

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


# 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)

Импортируйте библиотеку serverless_gpudistributed и модуль.
```
import serverless_gpu
from serverless_gpu import distributed
```

Оберните код обучения модели в функцию и декорируйте функцию декоратором @distributed.

@distributed(gpus=8, gpu_type='H100')
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-parellel 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()

Выполните распределенное обучение, вызвав распределенную функцию с определяемыми пользователем аргументами.
```
run_train.distributed(num_epochs=3, batch_size=1)
```
При выполнении в выходных данных ячейки блокнота появляется ссылка на запуск MLflow. Щелкните ссылку запуска MLflow или найдите ее на панели "Эксперимент ", чтобы просмотреть результаты выполнения. Дополнительные сведения о настройке имен экспериментов, отслеживания метрик и возобновлении выполнения см. в разделе "Отслеживание экспериментов" и "Наблюдаемость".

Сведения о распределённом выполнении

Бессерверный API GPU состоит из нескольких ключевых компонентов:

Менеджер вычислений: распределяет и управляет ресурсами
Среда выполнения. Управление средами и зависимостями Python
Загрузчик: управление выполнением и мониторингом заданий

При выполнении в распределенном режиме:

Функция сериализуется и распределяется по указанному количеству GPU
Каждый GPU запускает копию функции с одинаковыми параметрами
Среда синхронизирована на всех ГПУ
Результаты собираются и возвращаются из всех GPU
Управление жизненным циклом: распределенное выполнение выполняется в течение жизненного цикла записной книжки. Когда работа блокнота завершается, выполнение также прекращается. Декоратор @distributed имеет тайм-аут по умолчанию 3 часа. Чтобы задать пользовательский тайм-аут, передайте timeout в секундах или timeout=None, чтобы отключить его. Установка времени ожидания доступна в среде GPU версии 5 и выше.

API поддерживает популярные библиотеки параллельного обучения, такие как Distributed Data Parallel (DDP), Полностью шардингованный параллелизм данных (FSDP), DeepSpeed.

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

Для получения советов и шаблонов, которые помогут сделать ваш учебный процесс более эффективным и устойчивым, ознакомьтесь с руководством по эффективности и устойчивости.

FAQs

Где должен размещаться код загрузки данных?

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

from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, gpu_type='H100')
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

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

Узнать больше

Справочник по API см. в документации Serverless GPU Python API.

Обратная связь

Были ли сведения на этой странице полезными?

Last updated on 2026-07-30