Formação distribuída em cadernos

Importante

Este recurso está em versão Beta.

O @distributed decorador do Databricks AI Runtime SDK (databricks-sdk-air) é a forma mais conveniente de executar treino distribuído a partir de um notebook Databricks. Decora a tua função de treino, chama-a, e o AI Runtime executa-a em todas as GPUs do nó a que o teu portátil está ligado. O mesmo código escala de uma única GPU para várias GPUs, sem necessidade de aprovisionar um cluster nem de configurar um inicializador distribuído.

Pontos-chave nesta página:

• O @distributed decorador executa uma função de treino em todas as GPUs do teu nó a partir de dentro de um caderno.
• Suporta PyTorch DDP, FSDP e DeepSpeed, e move código de GPU única para multi-GPU com alterações mínimas.
• Liga o teu portátil a um acelerador 8xH100 para treino multi-GPU completo. Por defeito, @distributed usa todas as GPUs do nó.

Note

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

Início Rápido

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

Decore o seu evento de treino com @distributed, e depois chame-o 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 o acelerador H100 8x, que provisiona um único nó com 8 GPUs. Por defeito, @distributed usa as oito GPUs. Defina num_accelerators para usar menos. O accelerator_type parâmetro também é opcional e detetado automaticamente a partir do acelerador ao qual o seu notebook está ligado.

Cada chamada a .distributed() cria uma execução do MLflow (ou uma execução secundária aninhada, se já existir uma execução ativa) e imprime uma ligação para a execução na saída da célula. Para um guia completo e executável, veja o exemplo completo.

Estruturas suportadas

A @distributed API integra-se com as principais bibliotecas de treino distribuídas:

• PyTorch Distributed Data Parallel (DDP): Paralelismo padrão de dados multi-GPU.
• Fully Sharded Data Parallel (FSDP): Treino eficiente em memória para modelos de grande dimensão.
• DeepSpeed: A biblioteca de otimização da Microsoft para treino de grandes modelos.

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

Como o @distributed decorador trabalha

Ao chamar uma função decorada com .distributed(), o AI Runtime trata dos aspetos técnicos que, de outra forma, configuraria manualmente com um iniciador distribuído:

• Serialização e fan-out: 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 do ambiente: O ambiente Python e as dependências são replicados em todos os níveis, pelo que todos os processos executam o mesmo código.
• Variáveis de ambiente de classificação: Variáveis padrão, como as LOCAL_RANK que são preenchidas para cada processo. Leia-as na sua função para colocar o modelo e os dados no dispositivo correto.
• Recolha de resultados: Os valores de retorno são recolhidos de todas as classificações e devolvidos ao chamador.
• Rastreio do MLflow: Cada chamada .distributed() cria uma execução do MLflow, ou uma execução filha aninhada, se já houver uma ativa, pelo que as métricas registadas pela sua função são associadas à mesma execução.
• Ciclo de vida e tempo limite: A execução distribuída ocorre durante o ciclo de vida do notebook. Ao terminar o notebook, a execução é terminada. O decorador tem um tempo limite predefinido de 3 horas. Introduza timeout em segundos para o alterar, ou timeout=None para o desativar. Os tempos limite personalizados requerem os ambientes GPU v5 ou superior.

A API baseia-se 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 já não recebe novas funcionalidades. A partir do ambiente GPU versão 6, o Databricks AI Runtime SDK (databricks-sdk-air) substitui-o e fornece APIs equivalentes, com argumentos renomeados, sob o databricks.air namespace. O pacote que usas depende da versão do ambiente da GPU:

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

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

	serverless_gpu (preterido)	SDK de Runtime de IA (databricks.air)
Versões do ambiente GPU	6 anos ou menos	6 anos ou mais
Importação	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argumento de contagem de GPUs	gpus	num_accelerators
Argumento de tipo de GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valores do tipo 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 os dados de carregamento do tempo de execução de IA.

Para migrar um notebook:

1. Ligue o portátil ao ambiente 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 de acelerador. Os antigos nomes de argumentos e valores de tipo de GPU ainda funcionam em databricks.air, mas geram um FutureWarning. O suporte terminará para a versão 7 do ambiente GPU.

Dica

Diz ao Genie Code para fazer isto por ti:

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

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

Vindo do TorchDistributor

Se executares treino distribuído usando TorchDistributor em clusters Spark e a tua carga de trabalho cabe num único nó, a @distributed API é a substituição recomendada no AI Runtime. O seu código de modelo DDP, FSDP ou DeepSpeed não muda.

Feature	@distributed API (Interface de Programação de Aplicações)	TorchDistribuidor
Infraestrutura	Totalmente serverless, sem gestão de clusters	Requer um cluster Spark com trabalhadores da GPU
Setup	Decorador único, configuração mínima	Requer configuração do cluster Spark e do TorchDistributor.
Suporte à estrutura	PyTorch DDP, FSDP, DeepSpeed	Principalmente PyTorch DDP
Carregamento de dados	No decorador, utiliza volumes do Unity Catalog (UCVolumeDataset para transmissão contínua de dados de ficheiros)	Via Spark ou sistema de ficheiros

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

• Substitua a chamada TorchDistributor(...).run(train_fn, ...) pelo decorador @distributed em train_fn e, em seguida, inicie com train_fn.distributed(...).
• Remova o cluster Spark e a configuração do trabalhador da GPU. Ligue o seu portátil a um acelerador 8xH100 em vez disso. @distributed usa as oito GPUs por defeito.
• Mover o carregamento de dados dentro da função decorada. Ver carregamento de dados.
• Mantém o código existente do modelo DDP, FSDP ou DeepSpeed. O decorador apoia os três.

O @distributed decorador está em Beta e corre num único nó (ver Limitações), por isso não substitui todas as cargas de trabalho do TorchDistributor. Mantém as cargas de trabalho que dependem da integração com o Spark no TorchDistributor. Para executar treino distribuído a partir da sua máquina local ou entre múltiplos nós, use os comandos da CLI do Databricks para AI Runtime. Veja Utilizar a CLI do Databricks com runtime de IA.

Exemplo completo

O exemplo seguinte treina um modelo de perceptrão multicamada (MLP) em 8 GPUs H100 a partir de um portátil.

1. Configura o teu modelo e define funções de utilidade.

   
   # 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. Importa o distributed decorador a partir do SDK de Runtime da IA.

   from databricks.air import distributed
   

3. Envolve o código de treino do modelo numa função e decora a função com o @distributed decorador. A função decorada é o ponto de entrada para a execução distribuída, por isso define toda a lógica de treino, carregamento de dados e inicialização do modelo 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 treino distribuído chamando a função distribuída com argumentos definidos pelo utilizador.

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

5. Quando é executado, é gerada uma ligação para uma execução no MLflow na saída da célula do notebook. Clique no link da execução do MLflow ou encontre-o no painel Experimentos para ver os resultados da execução. Para detalhes sobre como personalizar nomes de experiências, rastrear métricas e retomar execuções, consulte Rastreio e observabilidade de experiências.

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, pelo que 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 ficheiros armazenados em volumes do Unity Catalog, use UCVolumeDataset de databricks.air.data, que transmite ficheiros com cache local e os distribui automaticamente entre ranks e workers. Para criar um ponto de verificação do treino distribuído num volume, use UCVolumeWriter e UCVolumeReader. Veja Carregar dados no AI Runtime e Checkpoint com Distributed Checkpoint (DCP).

Limitations

• @distributed suporta apenas um nó: ele corre através das GPUs do nó ao qual o seu notebook está ligado. Para treinar em múltiplos nós, use a CLI AI Runtime Databricks. Veja Utilizar a CLI do Databricks com runtime de IA.
• Se definires accelerator_type, deve corresponder ao acelerador ao qual o teu portátil está ligado (por exemplo, "GPU_8xH100"). Uma incompatibilidade faz falhar a carga de trabalho. O parâmetro é opcional e deteta-se automaticamente quando é omitido.
• O AI Runtime suporta o ambiente GPU v4 e superiores. Os limites de tempo personalizados (parâmetro timeout) requerem o ambiente GPU v5 ou superior.
• Por predefinição, o decorador expira ao fim de 3 horas. Introduza timeout em segundos para o alterar, ou timeout=None para o desativar.
• A execução ocorre no ciclo de vida do notebook. Ao terminar o notebook, a execução é terminada.

Recursos adicionais

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