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.
| Método e assinatura | Descrição |
|---|---|
init(self, num_processes, local_mode, use_gpu) |
Crie uma instância do 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 da CLI se main for um caminho de arquivo. |
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.
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.
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.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, odevice_idao qual os modelos e dados estão vinculados é determinado por você:device_id = int(os.environ["LOCAL_RANK"])Inicie o 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 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:
- Importar o repositório: importe o repositório externo como um repositório do Databricks.
- Criar um novo notebook Inicialize um novo Notebook do Azure Databricks no repositório.
- Iniciar treinamento distribuído Em uma célula do notebook, chame
TorchDistributorda 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.
Blocos de anotações de exemplo
Os exemplos de notebook a seguir demonstram como executar o treinamento distribuído com o PyTorch.