Preparação distribuída com o TorchDistributor
Este artigo descreve como executar 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 o PyTorch em seus clusters Spark, então ele permite que você inicie trabalhos de treinamento do PyTorch como trabalhos do Spark. Sob o capô, ele inicializa o ambiente e os canais de comunicação entre os trabalhadores e utiliza o comando torch.distributed.run CLI para executar treinamento distribuído entre os nós de trabalho.
A API TorchDistributor suporta os métodos mostrados na tabela a seguir.
| Método e assinatura | Description |
|---|---|
init(self, num_processes, local_mode, use_gpu) |
Crie uma instância de TorchDistributor. |
run(self, main, *args) |
Executa o treinamento distribuído invocando main(**kwargs) se main é uma função e executa o comando torchrun main *args CLI se main for um caminho de arquivo. |
Requisitos
- Faísca 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 acontecer inteiramente a partir de um notebook em sua máquina local ou um Notebook Databricks, você só precisará fazer pequenas alterações para preparar seu código para treinamento distribuído.
Preparar código de nó único: Prepare e teste o código de nó único com o PyTorch, o PyTorch Lightning ou outras estruturas baseadas no PyTorch/PyTorch Lightning, como a API do HuggingFace Trainer.
Preparar 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 englobado em uma função de treinamento que você pode usar com o
TorchDistributor.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 que você evite erros comuns de decapagem. Além disso, o que osdevice_idmodelos e dados estão vinculados é determinado por:device_id = int(os.environ["LOCAL_RANK"])Lançar treinamento distribuído: Instancie o
TorchDistributorcom os parâmetros desejados e chame.run(*args)para iniciar o treinamento.
Segue-se um exemplo de código de formação:
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 o Azure Databricks fazendo o seguinte:
- Importar o repositório: importe o repositório externo como uma pasta Databricks Git.
- Criar um novo bloco de anotações Inicialize um novo Bloco de Anotações do Azure Databricks no repositório.
- Iniciar treinamento distribuído Em uma célula de bloco de anotações, chame
TorchDistributorda seguinte forma:
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)
Resolução de Problemas
Um erro comum para o fluxo de trabalho do bloco de anotações é que os objetos não podem ser encontrados ou capturados 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 que é chamada com TorchDistributor(...).run(<func>) e dentro de quaisquer outras funções definidas pelo usuário chamadas no método de treinamento.
Falha NCCL: ncclInternalError: Internal check failed.
Quando você encontra esse erro durante o treinamento de vários nós, ele geralmente indica um problema com a comunicação de rede entre GPUs. Esse problema surge quando NCCL (NVIDIA Collective Communications Library) não pode usar determinadas interfaces de rede para comunicação GPU.
Para resolver esse erro, adicione o seguinte trecho no código de treinamento para usar a interface de rede primária.
import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
Blocos de notas de exemplo
Os exemplos de bloco de anotações a seguir demonstram como realizar treinamento distribuído com o PyTorch.