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), note 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 l'utilizzo di 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 worker master e uno o più nodi worker. Questi ruoli hanno i seguenti comportamenti:

  • Worker master: la VM con ranking 0. Questo nodo configura connessioni tra i nodi nel cluster.
  • Worker: i nodi rimanenti nel cluster. Ogni nodo esegue una parte di 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 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 ogni worker da utilizzare nel campo trainingInput.workerConfig.imageUri. Consulta l'elenco dei container 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 sarà trainingInput.masterConfig.imageUri. Spesso ha senso 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 distribuito supportato 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 il corso torch.nn.parallel.DistributedDataParallel. Ad esempio, se hai creato un modulo PyTorch chiamato model nel codice, aggiungi la seguente riga:

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

Per ulteriori informazioni sulla configurazione dell'addestramento distribuito, leggi la guida all'addestramento distribuito della documentazione di PyTorch.

Backend di addestramento distribuiti

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 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 ciascun nodo:

  • WORLD_SIZE: il numero totale di nodi nel cluster. Questa variabile ha lo stesso valore su tutti i nodi.
  • RANK: un identificatore univoco per ciascun nodo. Sul worker principale, è impostato su 0. Su 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 tutti i nodi.
  • MASTER_PORT: la porta su cui comunica il nodo worker master. Questa variabile ha lo stesso valore su tutti i nodi.

PyTorch utilizza queste variabili di ambiente per inizializzare il cluster.

Passaggi successivi