Configurazione dell'addestramento distribuito per PyTorch

Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

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

Questa guida presuppone che utilizzi i 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 della struttura del cluster di addestramento

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

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

Per informazioni su come specificare il nodo worker master e i nodi worker per il cluster di addestramento, consulta la sezione Specificare i tipi di macchine o i 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, quindi specifica l'immagine di un container Docker per ogni worker da utilizzare nel campo trainingInput.workerConfig.imageUri. Consulta l'elenco di contenitori PyTorch predefiniti.

Se utilizzi il comando gcloud ai-platform jobs submit training per creare il 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 è opportuno utilizzare lo stesso container PyTorch predefinito su ogni nodo.

Aggiornamento del codice di addestramento

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

import torch

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

Sostituisci BACKEND con uno dei backend di addestramento distribuiti supportati descritti nella sezione seguente. L'argomento della 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 per utilizzare la classe torch.nn.parallel.DistributedDataParallel. Ad esempio, se nel tuo codice hai creato un modulo PyTorch denominato model, aggiungi la seguente riga:

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

Per scoprire di più sulla configurazione della formazione distribuita, leggi la guida alla formazione distribuita.

Backend di addestramento distribuito

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

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

Scopri le 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. Per il worker master, è impostato su 0. Per ogni worker è impostato un valore diverso da 1 a WORLD_SIZE - 1.
  • MASTER_ADDR: il nome host del nodo worker master. Questa variabile ha lo stesso valore su ogni nodo.
  • MASTER_PORT: la porta su cui comunica il nodo worker master. Questa variabile ha lo stesso valore su ogni nodo.

PyTorch utilizza queste variabili di ambiente per inizializzare il cluster.

Passaggi successivi