Treinamento distribuído com o TorchDistributor

Este artigo descreve como executar treinamento distribuído em modelos de ML do PyTorch usando TorchDistributor.

O TorchDistributor é um módulo de software livre no PySpark que ajuda o treinamento distribuído com o PyTorch em seus clusters Spark. Com ele, você pode iniciar trabalhos de treinamento do PyTorch como trabalhos do Spark. Em segundo plano, ele inicializa o ambiente e os canais de comunicação entre os trabalhadores e utiliza o comando torch.distributed.run da CLI para executar o treinamento distribuído entre os nós de trabalho.

A API TorchDistributor oferece suporte aos métodos mostrados na tabela a seguir.

Requisitos

• Spark 3.4
• Databricks Runtime 13.0 ML ou superior

Fluxo de trabalho de desenvolvimento para notebooks

Se o processo de criação e treinamento do modelo ocorrer inteiramente em um notebook em seu computador local ou em um Databricks Notebook, você só precisará fazer pequenas alterações para preparar seu código para treinamento distribuído.

1. Preparar código de nó único: Prepare e teste o código de nó único com PyTorch, PyTorch Lightning ou outras estruturas baseadas no PyTorch/PyTorch Lightning, como a API do Trainer HuggingFace.

2. Preparar o código para treinamento distribuído padrão: Você precisa converter seu treinamento de processo único em treinamento distribuído. Tenha esse código distribuído incluído em uma função de treinamento que você pode usar com o TorchDistributor.

3. Mover importações dentro da função de treinamento: Adicione as importações necessárias, como import torch, dentro da função de treinamento. Isso permite evitar erros comuns de pickling. Além disso, o device_id ao qual os modelos e dados estão vinculados é determinado por você:

   device_id = int(os.environ["LOCAL_RANK"])
   

4. Inicie o treinamento distribuído: Instancie o TorchDistributor com os parâmetros desejados e chame .run(*args) para iniciar o treinamento.

Veja a seguir um exemplo de código de treinamento:

from pyspark.ml.torch.distributor import TorchDistributor

def train(learning_rate, use_gpu):
  import torch
  import torch.distributed as dist
  import torch.nn.parallel.DistributedDataParallel as DDP
  from torch.utils.data import DistributedSampler, DataLoader

  backend = "nccl" if use_gpu else "gloo"
  dist.init_process_group(backend)
  device = int(os.environ["LOCAL_RANK"]) if use_gpu  else "cpu"
  model = DDP(createModel(), **kwargs)
  sampler = DistributedSampler(dataset)
  loader = DataLoader(dataset, sampler=sampler)

  output = train(model, loader, learning_rate)
  dist.cleanup()
  return output

distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train, 1e-3, True)

Migrar o treinamento de repositórios externos

Se você tiver um procedimento de treinamento distribuído existente armazenado em um repositório externo, poderá migrar facilmente para o Azure Databricks fazendo o seguinte:

1. Importar o repositório: importe o repositório externo como um repositório do Databricks.
2. Criar um novo notebook Inicialize um novo Notebook do Azure Databricks no repositório.
3. Iniciar treinamento distribuído Em uma célula do notebook, chame TorchDistributor da seguinte maneira:

from pyspark.ml.torch.distributor import TorchDistributor

train_file = "/path/to/train.py"
args = ["--learning_rate=0.001", "--batch_size=16"]
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train_file, *args)

Solução de problemas

Um erro comum para o fluxo de trabalho do notebook é que os objetos não podem ser encontrados ou em conserva ao executar o treinamento distribuído. Isso pode ocorrer quando as instruções de importação da biblioteca não são distribuídas para outros executores.

Para evitar esse problema, inclua todas as instruções de importação (por exemplo, import torch) ambas na parte superior da função de treinamento que é chamada com TorchDistributor(...).run(<func>) e dentro de qualquer outra função definida pelo usuário chamada no método de treinamento.

Falha de NCCL: ncclInternalError: Internal check failed.

Quando você encontra esse erro durante o treinamento de vários nós, normalmente indica um problema com a comunicação de rede entre GPUs. Esse problema ocorre quando a NCCL (Biblioteca de Comunicações Coletivas NVIDIA) não pode usar determinados adaptadores de rede para comunicação de GPU.

Para resolver esse erro, adicione o snippet a seguir no código de treinamento para usar o adaptador de rede primário.

import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"

Blocos de anotações de exemplo

Os exemplos de notebook a seguir demonstram como executar o treinamento distribuído com o PyTorch.

Treinamento distribuído de ponta a ponta no notebook do Databricks

Obter notebook

Ajuste fino distribuído de um bloco de anotações de modelo de Rosto Abraçando

Obter notebook

Treinamento distribuído em um notebook de arquivo PyTorch

Obter notebook

Treinamento distribuído usando o notebook PyTorch Lightning

Obter notebook