Configurazione dell'addestramento distribuito per PyTorch

Questo documento spiega come creare un job di addestramento PyTorch distribuito. Quando crei un job di addestramento distribuito, AI Platform Training esegue il codice su un cluster di istanze di macchine virtuali (VM), noto anche come nodi, con variabili di ambiente che supportano l'addestramento PyTorch distribuito. In questo modo, il job di addestramento può scalare per gestire una grande quantità di dati.

Questa guida presuppone che tu stia utilizzando container PyTorch predefiniti per l'addestramento, come descritto nella Guida introduttiva a PyTorch. L'adattamento del codice PyTorch per l'addestramento distribuito richiede modifiche minime.

Specifica la struttura del cluster di addestramento

Per l'addestramento distribuito di PyTorch, configura il job in modo da utilizzare un nodo master worker e uno o più nodi worker. Questi ruoli hanno i seguenti comportamenti:

  • Master worker:la VM con rank 0. Questo nodo configura le connessioni tra i nodi del cluster.
  • Worker:i nodi rimanenti del cluster. Ogni nodo esegue una parte dell'addestramento, come specificato dal codice dell'applicazione di addestramento.

Per scoprire come specificare il nodo worker master e i nodi worker per il cluster di addestramento, consulta Specifica dei tipi di macchine o dei livelli di scalabilità.

Specifica delle immagini container

Quando crei un job di addestramento, specifica l'immagine di un container Docker che il worker master deve utilizzare nel campo trainingInput.masterConfig.imageUri e specifica l'immagine di un container Docker per ciascun worker da utilizzare nel campo trainingInput.workerConfig.imageUri. Consulta l'elenco dei container predefiniti di PyTorch.

Se utilizzi il comando gcloud ai-platform jobs submit training per creare il tuo job di addestramento, puoi specificare questi campi con i flag --master-image-uri e --worker-image-uri.

Tuttavia, se non specifichi il campo trainingInput.workerConfig.imageUri, il valore predefinito è trainingInput.masterConfig.imageUri. Spesso è consigliabile utilizzare lo stesso container PyTorch predefinito su ogni nodo.

Aggiornamento del codice di addestramento

Nell'applicazione di addestramento, aggiungi il seguente codice per inizializzare il cluster di addestramento:

import torch

torch.distributed.init_process_group(
    backend='BACKEND',
    init_method='env://'
)

Sostituisci BACKEND con uno dei backend per l'addestramento distribuito supportati descritti nella sezione seguente. L'argomento parola chiave init_method='env://' indica a PyTorch di utilizzare le variabili di ambiente per inizializzare la comunicazione nel cluster. Scopri di più nella sezione Variabili di ambiente di questa guida.

Inoltre, aggiorna il codice di addestramento in modo da utilizzare la classe torch.nn.parallel.DistributedDataParallel. Ad esempio, se hai creato un modulo PyTorch denominato model nel codice, aggiungi la riga seguente:

model = torch.nn.parallel.DistributedDataParallel(model)

Per scoprire di più sulla configurazione dell'addestramento distribuito, consulta la guida all'addestramento distribuito della documentazione di PyTorch.

Backend per l'addestramento distribuito

AI Platform Training supporta i seguenti backend per l'addestramento distribuito di PyTorch:

  • gloo: consigliato per i job di addestramento della CPU
  • nccl: consigliato per i job di addestramento GPU

Scopri di più sulle differenze tra i backend.

Variabili di ambiente

Quando crei un job di addestramento PyTorch distribuito, AI Platform Training imposta le seguenti variabili di ambiente su ogni nodo:

  • WORLD_SIZE: il numero totale di nodi nel cluster. Questa variabile ha lo stesso valore su ogni nodo.
  • RANK: un identificatore univoco per ogni nodo. Sul worker principale, questo valore è impostato su 0. Su ogni worker, è impostato su un valore diverso da 1 a WORLD_SIZE - 1.
  • MASTER_ADDR: il nome host del nodo worker principale. Questa variabile ha lo stesso valore su ogni nodo.
  • MASTER_PORT: la porta su cui comunica il nodo worker principale. Questa variabile ha lo stesso valore in ogni nodo.

PyTorch utilizza queste variabili di ambiente per inizializzare il cluster.

Passaggi successivi