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 una un cluster di istanze di macchine virtuali (VM) chiamate anche nodi, con variabili di ambiente che supportano l'addestramento di PyTorch distribuito. In questo modo, il job di addestramento può essere scalato per gestire una grande quantità di dati.

Questa guida presuppone l'utilizzo di container PyTorch predefiniti per l'addestramento, come descritto nella Guida introduttiva all'utilizzo di 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:

  • Worker principale:la VM con ranking pari a 0. Questo nodo e configurare le connessioni tra i nodi nel cluster.
  • Worker: i nodi rimanenti del cluster. Ogni nodo esegue una parte come specificato dal codice dell'applicazione di addestramento.

Per scoprire come specificare il nodo worker master e i nodi worker per per il cluster di addestramento, leggi Specifica dei tipi di macchina o della scalabilità di archiviazione.

Specifica delle immagini container

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

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

Tuttavia, se non specifichi il campo trainingInput.workerConfig.imageUri, il suo valore predefinito è trainingInput.masterConfig.imageUri. it spesso ha senso usare lo stesso container PyTorch predefinito su ogni nodo.

Aggiornamento del codice di addestramento

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

import torch

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

Sostituisci BACKEND con uno degli strumenti di addestramento distribuiti supportati descritti nella sezione che segue. 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 pagina 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 disponi ha creato un modulo PyTorch denominato model nel tuo codice, aggiungi la riga seguente:

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

Per saperne di più sulla configurazione dell'addestramento distribuito, leggi l'articolo PyTorch della documentazione alle API addestramento.

Backend di addestramento distribuiti

AI Platform Training supporta i seguenti backend per PyTorch distribuito formazione:

  • 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 cluster.

Passaggi successivi