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 CPUnccl
: 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 su0
. Per ogni worker è impostato un valore diverso da1
aWORLD_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
- Per ulteriori informazioni sull'addestramento con PyTorch su AI Platform Training, segui il tutorial Introduzione a PyTorch.
- Per scoprire di più sull'addestramento distribuito di PyTorch in generale, leggi la documentazione sulla formazione distribuita.