Рабочая нагрузка с несколькими GPU

Это важно

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

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

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

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

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

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	С помощью Spark или файловой системы

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

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

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

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

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

Параметр gpu_type в @distributed должен соответствовать типу акселератора, к которому подключен блокнот. Например, @distributed(gpus=8, gpu_type='H100') требует, чтобы ноутбук был подключен к акселератору H100. Использование несогласованного типа акселератора (например, подключение к A10 при указании H100) приведет к сбою рабочей нагрузки.

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

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs and GPU type
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

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

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


# Define the model
import os
import torch
import torch.distributed as dist
import torch.nn as nn

def setup():
    dist.init_process_group("nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

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_gpu и распределенный модуль.
```
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

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

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

Справочную информацию по API можно найти в документации безсерверного GPU API для Python.

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

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

Last updated on 2026-03-21