Treinamento distribuído em Notebooks

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.

O decorador @distributed da Serverless GPU Python API é a maneira mais conveniente de executar o treinamento distribuído a partir de um notebook do Databricks. Decore sua função de treinamento, chame-a, e o AI Runtime a executa em todas as GPUs no nó ao qual seu notebook está conectado. O mesmo código é escalado de GPU única para múltiplas GPUs sem cluster para provisionamento e sem inicializador distribuído para configurar.

dica

Principais conclusões

• The @distributed decorator runs a treinamento function across every GPU on your node from inside a notebook.
• Ele oferece suporte ao PyTorch DDP, FSDP e DeepSpeed, e migra códigos de GPU única para múltiplas GPUs com alterações mínimas.
• Conecte seu notebook a um acelerador 8xH100 e defina gpus=8 para o treinamento completo com várias GPUs.

Início rápido

O pacote serverless_gpu é pré-instalado quando o seu notebook está conectado a uma GPU serverless. Decore sua função de treinamento com @distributed e, em seguida, chame-a com .distributed():

Python

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

Cada chamada a .distributed() cria uma execução do MLflow (ou uma execução secundária aninhada, caso uma já esteja ativa) e imprime um link de execução na saída da célula. Para obter um passo a passo completo e executável, consulte Exemplo completo.

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 para múltiplas GPUs.
• Fully Sharded Data Parallel (FSDP) : treinamento eficiente em memória para modelos grandes.
• DeepSpeed : biblioteca de otimização da Microsoft para treinamento de modelos grandes.

Para cenários de treinamento reais que usam cada biblioteca, consulte exemplos de Notebooks.

Como o decorador @distributed funciona

Quando você chama uma função decorada com .distributed(), o AI Runtime lida com a mecânica que, de outra forma, você configuraria manualmente com um inicializador distribuído:

• Serialization and fan-out : The function is serialized and launched on each of the gpus you request. Cada GPU executa uma cópia da função com os mesmos argumentos.
• Environment synchronization : The Python environment and dependencies are replicated across all ranks, so every process runs the same code.
• Variáveis de ambiente de classificação : Variáveis padrão, como LOCAL_RANK, são preenchidas para cada processo. Leia-os em sua função para posicionar o modelo e os dados no dispositivo correto.
• Coleta de resultados : os valores de retorno são coletados de todas as classificações e retornados ao chamador.
• MLflow acompanhamento : cada chamada de .distributed() cria uma execução do MLflow, ou uma execução filho aninhada se uma já estiver ativa, para que as métricas registradas da sua função acabem na mesma execução.
• Lifecycle e tempo limite : A execução distribuída ocorre dentro do ciclo de vida do notebook. Encerrar o notebook encerra a execução. O decorador tem um tempo limite default de 3 horas. Passe timeout em segundos para alterá-lo, ou timeout=None para desativá-lo. Tempos limite personalizados exigem o ambiente de GPU v5 e superior.

A API se baseia nas bibliotecas padrão do PyTorch: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) e DeepSpeed.

Coming from TorchDistributor

If you run distributed PyTorch on Spark today with TorchDistributor and your workload fits on a single node, the serverless_gpu @distributed API is the recommended replacement for new aprendizagem profunda workloads. It removes the Spark cluster and gives you the same code path from single-GPU to multi-GPU.

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	Dentro do decorador, usa volumes do Unity Catalog (UCVolumeDataset para transmissão de dados de arquivos)	Via Spark ou sistema de arquivos

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	Dentro do decorador, usa volumes do Unity Catalog (UCVolumeDataset para transmissão de dados de arquivos)	Via Spark ou sistema de arquivos

Para migrar uma workload de nó único:

• Substitua a chamada a TorchDistributor(...).run(train_fn, ...) pelo decorador @distributed em train_fn e, em seguida, inicie com train_fn.distributed(...).
• Remova a configuração do cluster do Spark e do worker de GPU. Conecte seu notebook a um acelerador 8xH100 e defina gpus=8 em vez disso.
• Mova o carregamento de dados para dentro da função decorada. Consulte Carregamento de dados.
• Mantenha o seu código de modelo DDP, FSDP ou DeepSpeed existente. O decorador é compatível com os três.

@distributed é executado em um único nó (consulte Limitações), portanto, não substitui todas as cargas de trabalho do TorchDistributor. Mantenha as cargas de trabalho que dependem da integração com o Spark no TorchDistributor. Para executar o treinamento distribuído a partir da sua máquina local ou em vários nós, use a CLI do AI Runtime, que está em Pré-visualização pública. Consulte a CLI do AI Runtime.

Exemplo completo

O exemplo a seguir treina um modelo de perceptron multicamadas (MLP) em 8 GPUs H100 a partir de um notebook.

1. 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():
       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)

2. Importar a biblioteca serverless_gpu e o módulo distributed.

   Python

   import serverless_gpu
   from serverless_gpu import distributed

3. 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 é o ponto de entrada para a execução distribuída, portanto, defina toda a lógica de treinamento, o carregamento de dados e a inicialização do modelo dentro dela.

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

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

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 encontre-o no painel Experimento para ver os resultados da execução. Para obter detalhes sobre como personalizar nomes de experimentos, acompanhamento de métricas e retomar execuções, consulte Acompanhamento de experimentos e observabilidade.

Carregamento de dados

Coloque o código de carregamento de dados dentro da função @distributed. Um dataset pode exceder o tamanho máximo permitido por pickle, portanto, gerá-lo ou carregá-lo dentro do decorador evita erros de serialização:

Python

from serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Para dados baseados em arquivo armazenados em volumes do Unity Catalog, use UCVolumeDataset de serverless_gpu.data, que faz transmissão de arquivos com cache local e os particiona em ranks e workers automaticamente. Para fazer checkpoint de treinamento distribuído em um volume, use UCVolumeWriter e UCVolumeReader. Consulte Carregar dados em AI Runtime e Criar pontos de verificação do modelo.

Limitations

• O treinamento distribuído é executado nas GPUs no nó único ao qual o seu notebook está conectado. Para o treinamento completo com várias GPUs, conecte-se a um acelerador 8xH100, que provisiona um nó com 8 GPUs, e defina gpus=8.
• Accelerator type must match. If you set gpu_type in @distributed, it must match the accelerator your notebook is connected to ("H100" or "A10"). A mismatch causes the workload to fail. The parameter is optional and auto-detected when omitted.
• O AI Runtime recomenda o ambiente de GPU versão 4 e acima. Os tempos limites personalizados (o parâmetro timeout) exigem o ambiente de GPU v5 ou acima.
• The decorator times out after 3 hours by default. Pass timeout in seconds to change it, or timeout=None to disable it.
• A execução ocorre dentro do ciclo de vida do notebook. Encerrar o notebook encerra a execução.

Saber mais

• Para o decorador @distributed, GPUType e as APIs Ray, consulte a documentação de referência da Serverless GPU Python API.
• Para ver padrões que tornam seu pipeline de treinamento mais eficiente e resiliente, consulte o guia de desempenho e resiliência.
• Para cenários de treinamento de ponta a ponta, consulte Notebooks examples.