Condividi tramite

Training distribuito con TorchDistributor

  • Articolo

Questo articolo descrive come eseguire il training distribuito nei modelli di apprendimento automatico PyTorch usando TorchDistributor.

TorchDistributor è un modulo open source in PySpark che consente agli utenti di eseguire il training distribuito con PyTorch nei cluster Spark, in modo da consentire l'avvio di processi di training PyTorch come processi Spark. Inizializza l'ambiente e i canali di comunicazione tra i ruoli di lavoro e usa il comando dell'interfaccia della riga di comando torch.distributed.run per eseguire il training distribuito tra i nodi di lavoro.

L'API TorchDistributor supporta i metodi illustrati nella tabella seguente.

Metodo e firma Descrizione
init(self, num_processes, local_mode, use_gpu) Creare un'istanza di TorchDistributor.
run(self, main, *args) Esegue il training distribuito richiamando main(**kwargs) se main è una funzione ed esegue il comando dell'interfaccia della riga di comando torchrun main *args se main è un percorso di file.

Requisiti

  • Spark 3.4
  • Databricks Runtime 13.0 ML o versione successiva

Flusso di lavoro di sviluppo per notebook

Se il processo di creazione e training del modello avviene interamente da un notebook nel computer locale o in un Notebook di Databricks, è necessario apportare modifiche minori per preparare il codice per il training distribuito.

  1. Preparare il codice a nodo singolo: preparare e testare il codice a nodo singolo con PyTorch, PyTorch Lightning o altri framework basati su PyTorch/PyTorch Lightning come l'API HuggingFace Trainer.

  2. Preparare il codice per il training distribuito standard: è necessario convertire il training di un singolo processo in training distribuito. Avere tutto questo codice distribuito incluso all'interno di una funzione di training che è possibile usare con TorchDistributor.

  3. Spostare le importazioni all'interno della funzione di training: aggiungere le importazioni necessarie, ad esempio import torch, all'interno della funzione di training. In questo modo è possibile evitare errori comuni di selezione. Inoltre, i device_id modelli e i dati a cui sono associati sono determinati da:

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

  4. Avviare training distribuito: Creare un’istanza di TorchDistributor con i parametri desiderati e chiamare .run(*args) per avviare il training.

Di seguito è riportato un esempio di codice di training:

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)

Eseguire la migrazione del training da repository esterni

Se si dispone di una procedura memorizzata di training distribuito esistente in un repository esterno, è possibile eseguire facilmente la migrazione ad Azure Databricks eseguendo le operazioni seguenti:

  1. Importare il repository: importare il repository esterno come cartella Git di Databricks.
  2. Creare un nuovo notebook Inizializzare un nuovo Notebook di Azure Databricks all'interno del repository.
  3. Avviare il training distribuito In una cella del notebook chiamare TorchDistributor come segue:
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)

Risoluzione dei problemi

Un errore comune per il flusso di lavoro del notebook è che gli oggetti non possono essere trovati o prelevati durante l'esecuzione del training distribuito. Ciò può verificarsi quando le istruzioni di importazione della libreria non vengono distribuite ad altri executor.

Per evitare questo problema, includere tutte le istruzioni import (ad esempio, import torch) sia nella parte superiore della funzione di training chiamata con TorchDistributor(...).run(<func>) che all'interno di qualsiasi altra funzione definita dall'utente chiamata nel metodo di training.

Errore NCCL: ncclInternalError: Internal check failed.

Quando si verifica questo errore durante il training a più nodi, in genere indica un problema con la comunicazione di rete tra GPU. Questo problema si verifica quando NCCL (NVIDIA Collective Communications Library) non può usare determinate interfacce di rete per la comunicazione GPU.

Per risolvere questo errore, aggiungere il seguente frammento di codice di training per usare l'interfaccia di rete primaria.

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

Notebook di esempio

Gli esempi di notebook seguenti illustrano come eseguire il training distribuito con PyTorch.

Training distribuito end-to-end nel notebook di Databricks

Ottenere il notebook

Ottimizzazione distribuita di un notebook del modello Hugging Face

Ottenere il notebook

Training distribuito in un notebook di File PyTorch

Ottenere il notebook

Training distribuito con il notebook PyTorch Lightning

Ottenere il notebook