Compartilhar via

Carga de trabalho de várias GPUs

Importante

Esse recurso está em Beta. Os administradores do workspace podem controlar o acesso a esse recurso na página Visualizações . Consulte Gerenciar visualizações do Azure Databricks.

Você pode iniciar cargas de trabalho distribuídas em várias GPUs em um único nó usando a API Python de GPU sem servidor. A API fornece uma interface simples e unificada que abstrai os detalhes do provisionamento de GPU, da configuração do ambiente e da distribuição da carga de trabalho. Com alterações mínimas de código, você pode mover perfeitamente do treinamento de GPU única para a execução distribuída de várias GPUs do mesmo notebook.

Estruturas com suporte

A @distributed API se integra às principais bibliotecas de treinamento distribuídas:

  • DDP (PyTorch Distributed Data Parallel) – Paralelismo de dados de várias GPUs padrão.
  • FSDP (Fully Sharded Data Parallel) – treinamento eficiente em termos de memória para modelos grandes.
  • DeepSpeed — biblioteca de otimização da Microsoft para treinamento de modelo grande.

serverless_gpu API vs. TorchDistributor

A tabela a seguir compara a serverless_gpu@distributed API com TorchDistributor:

Característica serverless_gpu @distributed API Distribuidor de Tochas
Infraestrutura Totalmente sem servidor, sem gerenciamento de cluster Requer um cluster Spark com trabalhadores de GPU
Configuração Decorador único, configuração mínima Requer a configuração do cluster Spark e do TorchDistributor
Suporte à estrutura PyTorch DDP, FSDP, DeepSpeed Principalmente PyTorch DDP
Carregamento de dados Dentro do decorador, utiliza Unity Catalog Volumes Via Spark ou sistema de arquivos

A serverless_gpu API é a abordagem recomendada para novas cargas de trabalho de aprendizado profundo no Databricks. TorchDistributor permanece disponível para cargas de trabalho fortemente acopladas a clusters Spark.

Início rápido

A API de GPU sem servidor para treinamento distribuído é pré-instalada quando você está conectado a uma GPU sem servidor em notebooks e trabalhos do Databricks. Recomendamos o ambiente de GPU 4 e superior. Para usá-lo para treinamento distribuído, importe e use o decorador distributed para distribuir sua função de treinamento.

Coloque o código de treinamento do modelo em uma função e decore a função com o decorador @distributed. A função decorada torna-se o ponto de entrada para execução distribuída– toda a lógica de treinamento, o carregamento de dados e a inicialização do modelo devem ser definidos dentro dessa função.

Aviso

O parâmetro gpu_type em @distributed deve corresponder ao tipo de acelerador ao qual o notebook está conectado. Por exemplo, @distributed(gpus=8, gpu_type='H100') requer que seu notebook esteja conectado a um acelerador H100. O uso de um tipo de acelerador incompatível (como conectar-se ao A10 ao especificar H100) fará com que a carga de trabalho falhe.

O snippet de código abaixo mostra o uso básico 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():
    ...

Veja abaixo um exemplo completo que treina um modelo de MLP (perceptron de várias camadas) em 8 GPUs H100 de um notebook.

  1. Configure seu modelo e defina funções utilitárias.

    
# 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. Importe a biblioteca serverless_gpu e o módulo distribuído .

    import serverless_gpu
from serverless_gpu import distributed

  3. Coloque o código de treinamento do modelo em uma função e decore a função com o decorador @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()

  4. Execute o treinamento distribuído chamando a função distribuída com argumentos definidos pelo usuário.

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

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

Detalhes da execução distribuída

A API de GPU sem servidor consiste em vários componentes principais:

  • Gerenciador de computação: manipula a alocação e o gerenciamento de recursos
  • Ambiente de runtime: gerencia ambientes e dependências do Python
  • Inicializador: Orquestra a execução e o monitoramento da tarefa

Ao executar no modo distribuído:

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

A API dá suporte a bibliotecas de treinamento paralelas populares, como DDP ( Distributed Data Parallel ), FSDP ( Fully Sharded Data Parallel ), DeepSpeed.

Você pode encontrar cenários de treinamento distribuídos mais reais usando as várias bibliotecas em exemplos de notebook.

FAQs

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

Ao usar a API de GPU sem servidor para treinamento distribuído, mova o código de carregamento de dados no decorador @distributed. O tamanho do conjunto de dados pode exceder o tamanho máximo permitido pelo pickle, assim, é recomendável gerar o conjunto de dados dentro do decorador, conforme 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)
  ....

Saiba mais

Para a referência de API, consulte a documentação da API Python de GPU sem servidor .

  • Last updated on