Pular para o conteúdo principal
Última atualização em

Treinamento distribuído em Notebooks

info

Beta

Esse recurso está em Beta.

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.

Pontos-key nesta página:

  • 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.
nota

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

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

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

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 reais de treinamento que usam cada biblioteca, consulte exemplos de treinamento distribuído.

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 Usar a CLI do Databricks com o 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)
    ...

For file-based data stored in Unity Catalog volumes, use UCVolumeDataset from databricks.air.data, which transmissão files with local caching and partitions them across ranks and workers automatically. Para salvar checkpoints de treinamento distribuído em um volume, use UCVolumeWriter e UCVolumeReader. Consulte Carregar dados no AI Runtime e Criar checkpoint com Distributed Checkpoint (DCP).

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.

Recursos adicionais​

Nesta página