Treinamento distribuído com TorchDistributor

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

TorchDistributor é um módulo de código aberto no PySpark que ajuda os usuários a fazer treinamento distribuído com PyTorch em seus clusters Spark, permitindo que você inicie Job de treinamento PyTorch como Spark Job. Sob o capô, ele inicializa o ambiente e o canal de comunicação entre o worker e utiliza o comando CLI torch.distributed.run para executar o treinamento distribuído entre os nós worker .

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

Método e assinatura

Descrição

init(self, num_processes, local_mode, use_gpu)

Crie uma instância de TorchDistributor.

run(self, main, *args)

treinamento distribuído de execução invocando main(**kwargs) se main for uma função e executando o comando CLI torchrun main *args se main for um caminho de arquivo.

Requisitos

  • Spark 3.4

  • Databricks Runtime 13.0 ML ouacima

Desenvolvimento fluxo de trabalho para Notebook

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

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

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

  3. Mova as importações na função de treinamento: adicione as importações necessárias, como import torch, na função de treinamento. Isso permite evitar erros comuns de decapagem. Além disso, o device_id ao qual os modelos e dados estão vinculados é determinado por:

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

  4. Iniciar 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 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 Databricks fazendo o seguinte:

  1. Importar o repositório: Importar o repositório externo como uma pasta Git do Databricks.

  2. Criar um novo Notebook Inicialize um novo Databricks Notebook dentro do repositório.

  3. Inicie o treinamento distribuído Em uma célula Notebook , chame TorchDistributor como a seguir:

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 Notebook é que os objetos não podem ser encontrados ou selecionados durante a execução do treinamento distribuído. Isso pode acontecer 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) na parte superior da função de treinamento chamada com TorchDistributor(...).run(<func>) e dentro de qualquer outra função definida pelo usuário chamada no método de treinamento.

Falha na CUDA: peer access is not supported between these two devices

Este é um erro potencial no conjunto G5 de GPUs na AWS. Para resolver esse erro, adicione o seguinte snippet ao seu código de treinamento:

import os
os.environ["NCCL_P2P_DISABLE"] = "1"

Falha na NCCL: ncclInternalError: Internal check failed.

Quando o senhor encontra esse erro durante o treinamento de vários nós, ele normalmente indica um problema com a comunicação de rede entre as GPUs. Esse problema surge quando a NCCL (NVIDIA Collective Communications biblioteca) não pode usar determinadas interfaces de rede para comunicação com a GPU.

Para solucionar esse erro, adicione o seguinte trecho ao seu código de treinamento para usar a interface de rede primária.

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

Notebook de exemplo

Os exemplos Notebook a seguir demonstram como realizar treinamento distribuído com PyTorch.

Treinamento distribuído de ponta a ponta em Databricks Notebook

Abra o bloco de anotações em outra guia

Ajuste fino distribuído de um Notebookmodelo Hugging Face

Abra o bloco de anotações em outra guia

Treinamento distribuído em um Notebook de arquivo PyTorch

Abra o bloco de anotações em outra guia

Treinamento distribuído usando PyTorch Lightning Notebook

Abra o bloco de anotações em outra guia