carga de trabalho multi-GPU

info

Beta

Este recurso está em versão Beta. Os administradores do espaço de trabalho podem controlar o acesso a este recurso na página de Pré-visualizações . Veja as prévias do Gerenciador Databricks.

Você pode executar cargas de trabalho distribuídas em várias GPUs em um único nó usando a API Python para GPUs sem servidor. A API fornece uma interface simples e unificada que abstrai os detalhes do provisionamento de GPUs. Configuração do ambiente e distribuição da carga de trabalho. Com alterações mínimas no código, você pode migrar sem problemas. Do treinamento com uma única GPU à execução distribuída com múltiplas GPUs a partir do mesmo Notebook.

Estruturas suportadas

A API @distributed integra-se com as principais bibliotecas de treinamento distribuído:

PyTorch Distributed Data Parallel (DDP) — Paralelismo de dados padrão em múltiplas GPUs.
Paralelismo de Dados Totalmente Fragmentados (FSDP) — Treinamento com uso eficiente de memória para modelos grandes.
DeepSpeed — Biblioteca de otimização da Microsoft para treinamento de modelos de grande porte.

API serverless_gpu vs. TorchDistributor

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

Recurso	`serverless_gpu` `@distributed` API	TorchDistributor
Infraestrutura	Totalmente serverless, sem gerenciamento cluster	Requer um cluster Spark com um worker de GPU.
Configuração	Decorador único, configuração mínima	Requer a configuração de um cluster Spark e do TorchDistributor.
Suporte de estrutura	PyTorch DDP, FSDP, DeepSpeed	Principalmente PyTorch DDP
Carregamento de dados	Decorador interno, utiliza volumes Unity Catalog	Via Spark ou sistema de arquivos

A API serverless_gpu é a abordagem recomendada para novas cargas de trabalho de aprendizagem profunda no Databricks. O TorchDistributor continua disponível para cargas de trabalho fortemente acopladas a clusters Spark.

Comece rápido

A API de GPU serverless para treinamento distribuído é pré-instalada quando você se conecta a uma GPU serverless no Databricks Notebook e Job. Recomendamos o ambiente de GPU 4 ou superior. Para utilizá-lo em treinamento distribuído, importe e utilize o distributed Decorador para distribuir sua função de treinamento.

Envolva 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 a execução distribuída — toda a lógica de treinamento, carregamento de dados e inicialização do modelo devem ser definidos dentro dessa função.

atenção

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

O trecho de código abaixo mostra o uso básico de @distributed:

Python
# 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 ensina um modelo de perceptron multicamadas (MLP) em 8 GPUs H100 a partir de um Caderno:

Configure seu modelo e defina as funções de utilidade.

Python

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

Importe a biblioteca serverless_gpu e o módulo distribuído .
Python
```
import serverless_gpu
from serverless_gpu import distributed
```

Envolva o código de treinamento do modelo em uma função e decore a função com o decorador @distributed .

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

Execute o treinamento distribuído chamando a função distribuída com argumentos definidos pelo usuário.
Python
```
run_train.distributed(num_epochs=3, batch_size=1)
```
Ao ser executado, um link de execução MLflow será gerado na saída da célula do Notebook. Clique no link de execução do MLflow ou encontre-o no painel de Experimentos para ver os resultados da execução.

Detalhes da execução distribuída

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

Gerenciador de computação: Gerencia a alocação e o gerenciamento de recursos.
Ambiente Runtime : gerenciar ambientes e dependências Python
Iniciador: Orquestra a execução e o monitoramento de tarefas.

Ao executar em modo distribuído:

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

A API suporta bibliotecas populares de processamento paralelo, como Distributed Data Parallel (DDP). Dados Paralelos Totalmente Fragmentados (FSDP), DeepSpeed.

Você pode encontrar mais cenários reais de treinamento distribuído usando as várias bibliotecas nos exemplos de notebooks.

Perguntas frequentes

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

Ao usar a APIde GPU sem servidor Para treinamento distribuído, mova o código de carregamento de dados para dentro do decorador @distributed . O dataset O tamanho pode exceder o tamanho máximo permitido pelo pickle, portanto, recomenda-se gerar o dataset dentro do decorator, conforme mostrado abaixo:

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

Saber mais

Para obter informações sobre a API , consulte a documentação API Python para GPUs sem servidor .

Estruturas suportadas​

API serverless_gpu vs. TorchDistributor​

Comece rápido​

Detalhes da execução distribuída​

Perguntas frequentes​

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

Saber mais​