Partilhar via

Carga de trabalho multi-GPU

Importante

Este recurso está em versão Beta. Os administradores do espaço de trabalho podem controlar o acesso a esse recurso na página Visualizações . Consulte Gerenciar visualizações do Azure Databricks.

Pode lançar cargas de trabalho distribuídas entre várias GPUs num único nó usando a API Python da GPU Serverless. A API fornece uma interface simples e unificada que abstrai os detalhes do provisionamento da GPU, da configuração do ambiente e da distribuição da carga de trabalho. Com alterações mínimas de código, pode transitar sem problemas do treino com uma só GPU para a execução multi-GPU distribuída a partir do mesmo notebook.

Estruturas suportadas

A @distributed API integra-se com as principais bibliotecas de treino distribuídas:

  • PyTorch Distributed Data Parallel (DDP) — Paralelismo padrão de dados multi-GPU.
  • Fully Sharded Data Parallel (FSDP) — Treino eficiente em termos de memória para modelos grandes.
  • DeepSpeed — a biblioteca de otimização da Microsoft para treino de grandes modelos.

serverless_gpu API vs. TorchDistributor

A tabela seguinte compara a serverless_gpu@distributed API com o TorchDistributor:

Feature serverless_gpu @distributed API TorchDistribuidor
Infraestrutura Totalmente serverless, sem gestão de clusters Requer um cluster Spark com trabalhadores da GPU
Configuração Decorador único, configuração mínima Requer configuração do cluster Spark e do TorchDistributor.
Suporte à estrutura PyTorch DDP, FSDP, DeepSpeed Principalmente PyTorch DDP
Carregamento de dados Decorador de interiores, usa volumes do catálogo Unity Via Spark ou sistema de ficheiros

A serverless_gpu API é a abordagem recomendada para novas cargas de trabalho de deep learning em Databricks. O TorchDistributor mantém-se disponível para cargas de trabalho fortemente acopladas aos clusters Spark.

Início rápido

A API da GPU serverless para treino distribuído está pré-instalada quando está ligado a uma GPU serverless dentro dos notebooks e jobs do Databricks. Recomendamos o ambiente GPU 4 e superiores. Para usá-lo para treino distribuído, importe e utilize o decorador distributed para distribuir a sua função de treino.

Envolve o código de treino do modelo numa função e decora a função com o @distributed decorador. A função decorada torna-se o ponto de entrada para a execução distribuída — toda a lógica de treino, carregamento de dados e inicialização do modelo devem estar definidos dentro desta função.

Advertência

O parâmetro gpu_type em @distributed tem de corresponder ao tipo de acelerador ao qual o seu notebook está ligado. Por exemplo, @distributed(gpus=8, gpu_type='H100') requer que o seu portátil esteja ligado a um acelerador H100. Usar um tipo de acelerador incompatível (tal como ligar ao A10 enquanto se especifica o H100) fará o trabalho falhar.

O excerto de código abaixo mostra a utilização básica de @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():
    ...

Abaixo está um exemplo completo que treina um modelo de perceptrão multicamada (MLP) em 8 GPUs H100 a partir de um portátil:

  1. Configura o teu modelo e define funções de utilidade.

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

  2. Importa a biblioteca serverless_gpu e o módulo `distributed`.

    import serverless_gpu
from serverless_gpu import distributed

  3. Envolve o código de treino do modelo numa função e decora a função com o @distributed decorador.

    @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()

  4. Execute o treino distribuído chamando a função distribuída com argumentos definidos pelo utilizador.

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

  5. Quando for executada, é gerado um link de execução do MLflow na saída da célula do notebook. Clique no link da execução do MLflow ou encontre-o no painel Experimentos para ver os resultados da execução.

Detalhes de execução distribuída

A API de GPU serverless consiste em vários componentes-chave:

  • Gestor de computação: Gere a alocação e gestão de recursos
  • Ambiente de execução: Gere ambientes e dependências Python
  • Lançador: Orquestra a execução e monitorização do trabalho

Quando executar em modo distribuído:

  • A função é serializada e distribuída pelo número especificado de GPUs
  • Cada GPU executa uma cópia da função com os mesmos parâmetros
  • O ambiente está sincronizado em todas as GPUs
  • Os resultados são recolhidos e devolvidos de todas as GPUs

A API suporta bibliotecas de treino paralelas populares, como Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

Pode encontrar mais cenários reais de treino distribuído usando as várias bibliotecas em exemplos de cadernos.

FAQs

Onde deve ser colocado o código de carregamento dos dados?

Ao usar a API de GPU Serverless para treinamento distribuído, mova o código de carregamento de dados para dentro do decorador @distributed. O tamanho do conjunto de dados pode exceder o tamanho máximo permitido pelo pickle, por isso é recomendado criar o conjunto de dados dentro do decorador, como mostrado abaixo:

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

Mais informações

Para referência à API, consulte a documentação da API Python para GPU Serverless .

  • Last updated on