Google
Cloud Platform
Amazon Web Services
Microsoft AzureTreinamento 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 |
|---|---|
|
Crie uma instância de TorchDistributor. |
|
treinamento distribuído de execução invocando |
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.
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.
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.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, odevice_idao qual os modelos e dados estão vinculados é determinado por:device_id = int(os.environ["LOCAL_RANK"])
Iniciar treinamento distribuído: instancie o
TorchDistributorcom 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:
Importar o repositório: Importar o repositório externo como uma pasta Git do Databricks.
Criar um novo Notebook Inicialize um novo Databricks Notebook dentro do repositório.
Inicie o treinamento distribuído Em uma célula Notebook , chame
TorchDistributorcomo 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.
