Treinamento distribuído em notebooks

Importante

Esse recurso está em Beta.

O @distributed decorador do Databricks AI Runtime SDK (databricks-sdk-air) é a forma mais conveniente de executar treinamento distribuído a partir de um notebook Databricks. Decore sua função de treinamento, chame-a, e o AI Runtime roda em todas as GPUs do nó ao qual seu notebook está conectado. O mesmo código escala de GPU única para multiGPU, sem cluster para provisionar e sem launcher distribuído para configurar.

Pontos-chave desta página:

• O decorador @distributed executa uma função de treinamento em cada GPU do seu nó de dentro de um notebook.
• Ele suporta PyTorch DDP, FSDP e DeepSpeed, e move código de GPU única para multi-GPU com mudanças mínimas.
• Conecte seu notebook a um acelerador 8xH100 para treinamento completo multi-GPU. Por padrão, @distributed usa todas as GPUs do nó.

Observação

Esta página cobre o treinamento distribuído a partir de notebooks Databricks com o Databricks AI Runtime SDK. Para enviar cargas de trabalho de treinamento distribuídas da sua máquina local, use os comandos da CLI do Databricks para AI Runtime. Veja Usar a CLI do Databricks com o AI Runtime.

Início Rápido

Os exemplos nesta página utilizam o Databricks AI Runtime SDK, que é pré-instalado no ambiente GPU versão 6 e superior. Na versão 5 e abaixo, use o serverless_gpu pacote em vez disso. Consulte Migrar do serverless_gpu. Para conectar um notebook a um acelerador, veja Passo 1: Conecte um notebook a uma GPU.

Decore sua função de treinamento com @distributed e depois chame-a com .distributed():

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
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()

O treinamento distribuído requer um acelerador 8xH100, que provisiona um único nó com 8 GPUs. Por padrão, @distributed usa todas as oito GPUs. Defina num_accelerators para usar menos. O accelerator_type parâmetro também é opcional e detectado automaticamente pelo acelerador ao qual seu notebook está conectado.

Cada chamada de .distributed() cria uma execução no MLflow (ou uma execução filha aninhada, se já houver uma execução ativa) e exibe um link da execução na saída da célula. Para um guia completo e executável, veja o exemplo completo.

Estruturas com suporte

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

• DDP (Distributed Data Parallel) do PyTorch: paralelismo de dados de várias GPUs padrão.
• FSDP (Fully Sharded Data Parallel): treinamento com eficiência de memória para modelos grandes.
• DeepSpeed: biblioteca de otimização do Microsoft para treinamento de modelo grande.

Para cenários reais de treinamento que usam cada biblioteca, veja exemplos de treinamento distribuído.

Como o @distributed decorador trabalha

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

• Serialização e espalhamento: A função é serializada e lançada em cada GPU selecionada. Cada GPU executa uma cópia da função com os mesmos argumentos.
• Sincronização de ambientes: O ambiente Python e as dependências são replicados em todos os níveis, então todo processo executa o mesmo código.
• Variáveis de ambiente de rank: Variáveis padrão, como LOCAL_RANK, são definidas para cada processo. Leia-os em sua função para colocar o modelo e os dados no dispositivo correto.
• Coleta de resultados: Os valores de retorno são coletados de todas as classificações e devolvidos ao chamador.
• Rastreamento do MLflow: Cada chamada .distributed() cria uma execução do MLflow, ou uma execução filha aninhada, se já houver uma execução ativa, de modo que as métricas registradas por sua função sejam associadas à mesma execução.
• Ciclo de vida e timeout: A execução distribuída roda dentro do ciclo de vida do notebook. Encerrar o notebook encerra a execução. O decorador tem um timeout padrão de 3 horas. Passe timeout em segundos para alterá-lo ou timeout=None para desativá-lo. Limites de tempo personalizados exigem o ambiente GPU v5 ou superior.

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

Migrar de serverless_gpu

O serverless_gpu pacote está obsoleto e não recebe mais novos recursos. A partir da versão 6 do ambiente GPU, o Databricks AI Runtime SDK (databricks-sdk-air) o substitui e fornece APIs equivalentes, com argumentos renomeados, sob o namespace databricks.air. O pacote que você usa depende da versão do ambiente da GPU:

• Versão 5 ou inferior: Inclua apenas serverless_gpu. Continue usando serverless_gpu nessas versões.
• Versão 6: Inclui ambos os pacotes, mas apenas o SDK de Runtime de IA recebe atualizações. A versão 6 é a última versão do ambiente da GPU que inclui serverless_gpu.

O SDK de Runtime de IA é um pacote separado do SDK Databricks para Python (databricks-sdk).

	serverless_gpu (preterido)	SDK de Runtime de IA (databricks.air)
Versões do ambiente GPU	6 anos ou menos	6 ou mais
Importar	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argumento de contagem de GPUs	gpus	num_accelerators
Argumento do tipo de GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valores dos tipos de GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Módulo de carregamento de dados e checkpointing	serverless_gpu.data	databricks.air.data. Consulte Carregar dados no AI Runtime.

Para migrar um notebook:

1. Conecte o notebook ao ambiente da GPU versão 6 ou superior, que inclui o AI Runtime SDK.
2. Alterar importações de serverless_gpu para databricks.air.
3. Renomeie gpus para num_accelerators e gpu_type para accelerator_type, e use os novos valores do tipo acelerador. Os antigos nomes de argumentos e valores de tipo de GPU ainda funcionam em databricks.air, mas emitem um FutureWarning. O suporte será encerrado no ambiente de GPU versão 7.

Tip

Diga ao Genie Code para fazer isso por você:

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

Para cargas de trabalho da CLI do Databricks que usam databricks.air.data, adicione databricks-sdk-air[data] ao environment.dependencies da carga de trabalho.

Vindo da TorchDistributor

Se você roda treinamento distribuído usando TorchDistributor em clusters Spark e sua carga de trabalho cabe em um único nó, a @distributed API é a alternativa recomendada no AI Runtime. Seu código de modelo DDP, FSDP ou DeepSpeed não muda.

Feature	@distributed Interface de Programação de Aplicações (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 para estrutura	PyTorch DDP, FSDP, DeepSpeed	Principalmente PyTorch DDP
Carregamento de dados	Dentro do decorador, utiliza volumes do Catálogo do Unity (UCVolumeDataset para dados de arquivos de streaming)	Via Spark ou sistema de arquivos

Para migrar uma carga de trabalho de nó único:

• Substitua a chamada @distributed pelo decorador train_fn em train_fn.distributed(...), depois inicie com TorchDistributor(...).run(train_fn, ...).
• Remova o cluster do Spark e a configuração do trabalhador da GPU. Conecte seu notebook a um acelerador 8xH100 em vez disso. @distributed usa todas as oito GPUs por padrão.
• Mova o carregamento de dados dentro da função decorada. Veja Carregamento de dados.
• Mantenha seu código de modelo DDP, FSDP ou DeepSpeed existente. O decorador apoia os três.

O @distributed decorador está em Beta e roda em um único nó (veja Limitações), então ele não substitui todas as cargas de trabalho do TorchDistributor. Mantenha cargas de trabalho que dependem da integração com o Spark no TorchDistributor. Para rodar treinamento distribuído a partir da sua máquina local ou entre múltiplos nós, use os comandos da CLI do Databricks para o AI Runtime. Veja Usar a CLI do Databricks com o AI Runtime.

Exemplo completo

O exemplo a seguir treina um modelo de perceptron multicamada (MLP) em 8 GPUs H100 a partir 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():
       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. Importe o distributed decorador do SDK de Runtime da IA.

   from databricks.air import distributed
   

3. Encapsule o código de treinamento do modelo em uma função e aplique o decorador @distributed a essa função. A função decorada é o ponto de entrada para a execução distribuída, então defina toda a lógica de treinamento, carregamento de dados e inicialização de modelos dentro dela.

   @distributed()
   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-parallel 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. Ao ser executado, um link para a 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. Para obter detalhes sobre como personalizar nomes de experimentos, controlar métricas e retomar execuções, consulte Acompanhamento e observabilidade de experimentos.

Carregamento de dados

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

from databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Para dados baseados em arquivos armazenados em volumes do Catálogo do Unity, utilize UCVolumeDataset de databricks.air.data, que realiza o streaming de arquivos com cache local e os particiona automaticamente entre ranks e workers. Para salvar checkpoints de treinamento distribuído em um volume, utilize UCVolumeWriter e UCVolumeReader. Veja Carregar dados no AI Runtime e Checkpoint com Checkpoint Distribuído (DCP).

Limitations

• @distributed suporta apenas um nó: ele funciona através das GPUs do nó ao qual seu notebook está conectado. Para treinar em múltiplos nós, use a CLI do AI Runtime Databricks. Veja Usar a CLI do Databricks com o AI Runtime.
• Se você definir accelerator_type, ele deve corresponder ao acelerador ao qual seu notebook está conectado (por exemplo, "GPU_8xH100"). Uma incompatibilidade faz a carga de trabalho falhar. O parâmetro é opcional e é detectado automaticamente quando omitido.
• O AI Runtime suporta ambiente de GPU v4 e superiores. Timeouts personalizados (o parâmetro timeout) exigem ambientes de GPU v5 ou superior.
• O decorador expira após 3 horas por padrão. Passe timeout em segundos para alterá-lo ou timeout=None para desativá-lo.
• A execução é executada dentro do ciclo de vida do notebook. Encerrar o notebook encerra a execução.

Recursos adicionais

• Para o decorador @distributed, AcceleratorType e as APIs do Ray, veja a referência da API do SDK do Databricks AI Runtime.
• Para padrões que tornam seu pipeline de treinamento mais eficiente e resiliente, veja o guia de desempenho e resiliência.
• Para cenários de treinamento de ponta a ponta, veja exemplos de treinamento distribuído.