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.

nota

Esta página aborda o treinamento distribuído de Notebooks Databricks com a API Python Serverless de GPU. Para enviar cargas de trabalho de treinamento distribuído de sua máquina local, use a CLI do AI Runtime, que está em Pré-visualização pública. Consulte CLI do AI Runtime.

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.

nota

O treinamento distribuído requer um acelerador 8xH100, que provisiona um único nó com 8 GPUs. Ao usar o decorador @distributed, defina gpus=8. O parâmetro gpu_type é opcional e detectado automaticamente do acelerador ao qual seu Notebook está conectado.

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.

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	Dentro do decorador, usa volumes do Unity Catalog (`UCVolumeDataset` para dados de arquivo de transmissão)	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 dados de arquivo de transmissão)	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; portanto, toda a lógica de treinamento, carregamento de dados e inicialização do modelo deve ser definida dentro desta função.

Para iniciar a execução distribuída, chame sua função decorada usando train_function.distributed(). Cada chamada cria automaticamente uma execução de experimento do MLflow ou uma execução filha aninhada se uma já estiver ativa.

atenção

Ao definir gpu_type em @distributed, certifique-se de que ele corresponda ao tipo de acelerador ao qual seu Notebook está conectado ("H100" ou "A10"). A especificação do tipo de acelerador incorreto causará a falha da carga de trabalho.

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

Python
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs.
# gpu_type='H100' is optional and will be auto-detected if not set.
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

run_train.distributed()

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

Importar a biblioteca serverless_gpu e o módulo distributed.
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)
```
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.

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.
Gerenciamento do ciclo de vida: A execução distribuída ocorre dentro do ciclo de vida do Notebook. Quando o Notebook termina, a execução também termina. O decorador @distributed tem um tempo limite default de 3 horas. Para definir um tempo limite personalizado, passe timeout em segundos, ou timeout=None para desativá-lo. A configuração de um tempo limite está disponível no ambiente GPU v5 e acima.

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

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.

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​