Treinamento distribuído com TorchDistributor

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

TorchDistributor é um módulo de código aberto em PySpark que ajuda os usuários a fazer treinamento distribuído com PyTorch em seu clustering Spark, permitindo que o senhor inicie o PyTorch treinamento Job como Spark Job. Sob o capô, ele inicializa o ambiente e o canal de comunicação entre o trabalhador e utiliza o comando CLI torch.distributed.run para executar o treinamento distribuído nos 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 do TorchDistributor.
run(self, main, *args)	execução distribuiu o treinamento invocando main(**kwargs) se main for uma função e executando o CLI comando torchrun main *args se main for um caminho de arquivo.

Requisitos

• Spark 3.4
• Databricks Runtime 13.0 ML ouacima

Desenvolvimento de fluxo de trabalho para Notebook

Se a criação do modelo e o processo de treinamento ocorrerem inteiramente em um Notebook em sua máquina local ou em um Databricks Notebooko senhor só precisará fazer pequenas alterações para deixar o código pronto para o treinamento distribuído.

1. Preparar 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 o código para o treinamento distribuído padrão: O senhor precisa converter seu treinamento de processo único em treinamento distribuído. Tenha esse código distribuído todo englobado em uma função de treinamento que o senhor pode usar com o site TorchDistributor.

3. Mova as 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 que você evite erros comuns de decapagem. Além disso, o device_id ao qual os modelos e os dados estão vinculados é determinado por:

   Python

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

4. Iniciar o treinamento distribuído: Instanciar o site TorchDistributor com os parâmetros desejados e chamar .run(*args) para iniciar o treinamento.

A seguir, um exemplo de código de treinamento:

Python

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 o senhor tiver um procedimento de treinamento distribuído existente armazenado em um repositório externo, poderá migrar facilmente para a 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 Inicializar um novo Databricks Notebook dentro do repositório.
3. Iniciar treinamento distribuído Em uma célula do Notebook, chame TorchDistributor da seguinte forma:

Python

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 no fluxo de trabalho do Notebook é que os objetos não podem ser encontrados ou decapados 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 outro executor.

Para evitar esse problema, inclua todas as instruções de importação (por exemplo, import torch) 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 na CUDA: peer access is not supported between these two devices

Esse é um possível erro no conjunto de GPUs G5 no AWS. Para resolver esse erro, adicione o seguinte trecho em seu código de treinamento:

Python

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

Falha no NCCL: ncclInternalError: Internal check failed.

Quando o senhor encontra esse erro durante o treinamento com 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.

Python

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

Exemplo de notebook

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

Treinamento distribuído de ponta a ponta em Databricks Notebook

Open notebook in new tab

Ajuste fino distribuído de um modelo Hugging Face Notebook

Open notebook in new tab

Treinamento distribuído em um notebook PyTorch File Notebook

Open notebook in new tab

Treinamento distribuído usando o site PyTorch Lightning Notebook

Open notebook in new tab