Treinamento distribuído em notebooks

info

Esse recurso está na versão Beta.

O decorador @distributed do SDK do Databricks AI Runtime (databricks-sdk-air) é 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 realiza a execução dela em todas as GPUs no nó ao qual o Notebook está conectado. O mesmo código escala de GPU única para GPU múltipla sem clusters para provisionamento e nenhum lançador distribuído para configurar.

key pontos nesta página:

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

nota

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

Início rápido

Os exemplos nesta página usam o SDK do Databricks AI Runtime, que vem pré-instalado na versão 6 e acima do ambiente de GPU. Na versão 5 e abaixo, use o pacote serverless_gpu em vez disso. Consulte Migrar de serverless_gpu. Para conectar um Notebooks a um acelerador, consulte O passo 1: Anexar um Notebooks a uma GPU.

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

Python

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 ou 8xB300, que provisiona um único nó com 8 GPUs. Por default, @distributed usa todas as oito GPUs. Defina num_accelerators para usar menos. O parâmetro accelerator_type também é opcional e detectado automaticamente a partir do acelerador ao qual seu notebook está conectado.

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

Estruturas compatíveis

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

• PyTorch Distributed Data Parallel (DDP) : Paralelismo de dados multi-GPU padrão.
• Fully Sharded Data Parallel (FSDP) : treinamento eficiente em memória para grandes modelos.
• 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 treinamento distribuído.

Como funciona o decorador @distributed

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

• Serialização e distribuição : A função é serializada e iniciada em cada GPU selecionada. Cada GPU faz a execução de uma cópia da função com os mesmos argumentos.
• Sincronização de ambiente : O ambiente Python e as dependências são replicados em todas as hierarquias, então cada processo faz a execução do mesmo código.
• Variáveis de ambiente de classificação : variáveis padrão como LOCAL_RANK são preenchidas para cada processo. Leia-as em sua função para colocar o modelo e os dados no dispositivo correto.
• Coleta de resultados : valores de retorno são coletados de todos os ranks e retornados ao chamador.
• Acompanhamento do MLflow : cada chamada .distributed() cria uma execução do MLflow, ou uma execução filha aninhada se uma já estiver ativa, para que as métricas logadas de sua função caiam na mesma execução.
• Ciclo de vida e tempo limite : A execução distribuída é executada 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. Os timeouts personalizados exigem ambiente de GPU v5 e acima.

A API baseia-se nas bibliotecas padrão do PyTorch: Paralelismo de Dados Distribuídos (DDP), Paralelismo de Dados Totalmente Fragmentado (FSDP) e DeepSpeed.

Migrar de serverless_gpu

O pacote serverless_gpu foi descontinuado e não recebe mais novos recursos. A partir da versão 6 do ambiente GPU, o SDK de runtime do Databricks AI (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 GPU:

• Versão 5 e abaixo: inclua apenas serverless_gpu. Continue a usar serverless_gpu nessas versões.
• Versão 6: Inclui ambos os pacotes, mas apenas o SDK do AI Runtime recebe atualizações. A versão 6 é a última versão do ambiente GPU que inclui serverless_gpu.

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

	serverless_gpu (descontinuado)	SDK de AI Runtime (databricks.air)
Versões de ambiente de GPU	6 e abaixo	6 e acima
Importar	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argumento de contagem de GPU	gpus	num_accelerators
Argumento de tipo de GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valores de tipo de GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Módulo de carregamento de dados e verificação	serverless_gpu.data	databricks.air.data. Consulte Carregar dados no AI Runtime.

	serverless_gpu (descontinuado)	SDK de AI Runtime (databricks.air)
Versões de ambiente de GPU	6 e abaixo	6 e acima
Importar	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argumento de contagem de GPU	gpus	num_accelerators
Argumento de tipo de GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valores de tipo de GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Módulo de carregamento de dados e verificação	serverless_gpu.data	databricks.air.data. Consulte Carregar dados no AI Runtime.

Para migrar um notebook:

1. Conecte o notebook à versão 6 ou acima do ambiente de GPU, que inclui o AI Runtime SDK.
2. Altere as 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 de tipo de acelerador. Os nomes de argumentos antigos e os valores de tipo de GPU ainda funcionam em databricks.air, mas emitem um FutureWarning. O suporte terminará na versão 7 do ambiente de GPU.

prompt

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 do TorchDistributor

Se você fizer uma execução de treinamento distribuído usando TorchDistributor em clusters Spark e seu workload couber em um único nó, a API @distributed é a substituição recomendada no AI Runtime. Seu código de modelo DDP, FSDP ou DeepSpeed não muda.

Recurso	@distributed API	TorchDistributor
Infraestrutura	Totalmente serverless, sem gerenciamento de cluster	Requer um cluster Spark com worker de GPU
Configuração	Decorador único, configuração mínima	Requer clusters Spark e configuração do TorchDistributor
Suporte ao framework	PyTorch DDP, FSDP, DeepSpeed	Principalmente PyTorch DDP
Carregamento de dados	Dentro do decorador, usa volumes do Unity Catalog (UCVolumeDataset para dados de arquivo de transmissão)	Via Spark ou sistema de arquivos

Recurso	@distributed API	TorchDistributor
Infraestrutura	Totalmente serverless, sem gerenciamento de cluster	Requer um cluster Spark com worker de GPU
Configuração	Decorador único, configuração mínima	Requer clusters Spark e configuração do TorchDistributor
Suporte ao framework	PyTorch DDP, FSDP, DeepSpeed	Principalmente PyTorch DDP
Carregamento de dados	Dentro do decorador, usa volumes do Unity Catalog (UCVolumeDataset para dados de arquivo de transmissão)	Via Spark ou sistema de arquivos

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

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

O decorador @distributed está em Beta e sua execução ocorre em um único nó (consulte Limitações), por isso não substitui todas as workloads do TorchDistributor. Mantenha as workloads que dependem da integração Spark no TorchDistributor. Para a execução de treinamento distribuído a partir de sua máquina local ou em vários nós, use os comandos da CLI do Databricks para AI Runtime. Consulte Usar a CLI do Databricks com AI Runtime.

Exemplo completo

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

1. Configure seu modelo e defina 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. Importe o decorador distributed do SDK do AI Runtime.

   Python

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

   Python

   @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. A execução do 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 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, rastrear 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 função @distributed. Um dataset pode exceder o tamanho máximo permitido por pickle, então gerá-lo ou carregá-lo dentro do decorador evita erros de serialização:

Python

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 Unity Catalog, use UCVolumeDataset de databricks.air.data, que faz transmissão de arquivos com cache local e os particiona automaticamente entre ranques e workers. Para verificar o treinamento distribuído em um volume, use UCVolumeWriter e UCVolumeReader. Consulte Carregar dados no AI Runtime e Verificar com Distributed Checkpoint (DCP).

Limitações

• @distributed suporta apenas um único nó: ele entra em execução em todas as GPUs no nó ao qual seu notebook está conectado. Para ensinar em vários nós, use a CLI do AI Runtime Databricks. Consulte 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 causa a falha do workload. O parâmetro é opcional e detectado automaticamente quando omitido.
• O AI Runtime oferece suporte ao ambiente GPU v4 e acima. Tempos limite personalizados (o parâmetro timeout) exigem o ambiente GPU v5 e acima.
• O decorador expira após 3 horas por default. Passe timeout em segundos para alterá-lo, ou timeout=None para desativá-lo.
• A execução acontece dentro do ciclo de vida do notebook. Encerrar o notebook encerra a execução.

Recursos adicionais

• Para o decorador @distributed, AcceleratorType e APIs Ray, consulte a referência da API do SDK do Databricks AI Runtime.
• Para 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 exemplos de treinamento distribuído.