Addestra modelli di machine learning su larga scala su GKE con il checkpointing multilivello

Questa pagina mostra come utilizzare il checkpointing multi-livello per archiviare e gestire in modo affidabile i checkpoint durante l'addestramento del modello di machine learning su GKE. L'archiviazione e la gestione dei checkpoint sono fondamentali per i job di addestramento su larga scala, definiti come quelli che utilizzano più di migliaia di nodi. Le interruzioni di questi job su larga scala sono frequenti (potenzialmente orarie) e il ripristino può essere lento.

Vantaggi

L'utilizzo del checkpointing multilivello offre i seguenti vantaggi:

  • Gestione completamente orchestrata dei dati dei checkpoint, inclusi backup, replica e ripristino automatico per i seguenti carichi di lavoro:
  • Recupero rapido dei job di addestramento da un checkpoint archiviato nel nodo locale. Puoi anche eseguire il recupero utilizzando i checkpoint archiviati in un altro nodo del cluster di addestramento.
  • Ripristino rapido dei job di addestramento da un checkpoint archiviato in un backup di Cloud Storage negli scenari peggiori, in cui non sono presenti checkpoint nel cluster.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
    • Attiva l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo gcloud components update.

Requisiti

Il checkpointing multilivello richiede la versione 1.32.4-gke.1415000 o successive del cluster GKE.

Limitazioni

  • I cluster Autopilot non sono supportati.

Configurare i nodi GKE per l'utilizzo del checkpointing multilivello

Questa sezione spiega come configurare i nodi GKE su cluster nuovi ed esistenti.

Configura i nodi su un nuovo cluster

  1. Crea un cluster con il checkpoint multi-livello, il driver CSI di Cloud Storage FUSE e Workload Identity Federation for GKE abilitati. Se utilizzi le sezioni TPU per il tuo workload di machine learning, devi modificare il comando di creazione del cluster per includere la configurazione di un pool di nodi di sezioni TPU.

    gcloud container clusters create CLUSTER_NAME \
    --addons=HighScaleCheckpointing,GcsFuseCsiDriver  \
    --node-locations=NODE_LOCATION \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --cluster-version=CLUSTER_VERSION
    --location=CLUSTER_LOCATION \
    --machine-type=MACHINE_TYPE \
    --num-nodes=NUM_NODES

    Sostituisci i seguenti valori:

    • CLUSTER_NAME: il nome del cluster.
    • NODE_LOCATION: la zona per i nodi del cluster. Qui si trova la capacità TPU.
    • PROJECT_ID: il tuo Google Cloud ID progetto.
    • CLUSTER_VERSION: la versione del cluster. 1.32.4-gke.1415000 è la versione minima supportata.
    • CLUSTER_LOCATION: la regione in cui vuoi creare il cluster.
    • MACHINE_TYPE: il tipo di macchina utilizzato per i nodi che eseguono componenti come il controller JobSet e il controller di checkpointing multilivello. Per l'addestramento su larga scala, ti consigliamo di utilizzare almeno e2-standard-4 macchine. Non utilizzerai questo tipo di macchina per l'addestramento del modello. Al contrario, creerai pool di nodi separati per questo scopo, spesso utilizzando famiglie di VM ottimizzate per l'acceleratore.
    • NUM_NODES: il numero di nodi da creare in ciascuna delle zone del cluster.

Configurare i nodi su un cluster esistente

Per utilizzare il checkpointing multilivello con un cluster esistente, attivalo insieme al driver CSI di Cloud Storage FUSE e alla federazione delle identità per i carichi di lavoro per GKE con i seguenti comandi. La versione del cluster esistente deve essere successiva alla versione 1.32.3-gke.1170000.

  1. Abilita la federazione delle identità per i workload per GKE:

    gcloud container clusters update CLUSTER_NAME \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --location=CLUSTER_LOCATION

    Sostituisci i seguenti valori:

    • CLUSTER_NAME: il nome del cluster.
    • PROJECT_ID: il tuo Google Cloud ID progetto.
    • CLUSTER_LOCATION: la regione del cluster.

  2. Abilita il checkpointing multilivello e il driver CSI di Cloud Storage FUSE:

    gcloud container clusters update CLUSTER_NAME \
    --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
    --location=CLUSTER_LOCATION

Configura le autorizzazioni per utilizzare il checkpointing multilivello

Questa sezione illustra come configurare le autorizzazioni per utilizzare il checkpointing multilivello.

Concedere l'accesso ai bucket Cloud Storage

I volumi effimeri utilizzati dal driver CSI Multi-Tier Checkpointing devono utilizzare bucket Cloud Storage esistenti.

Per archiviare i checkpoint in un bucket Cloud Storage, il checkpointing multilivello deve accedere al bucket. Concedi il ruolo IAM Storage Object User (roles/storage.objectUser) sul bucket al service account Kubernetes per il checkpoint multilivello.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
    --role "roles/storage.objectUser"

Sostituisci i seguenti valori:

(Facoltativo) Concedi l'accesso all'account di servizio Compute Engine predefinito

Se le tue istanze Compute Engine hanno bisogno dell'accesso in lettura al bucket Cloud Storage, concedi il ruolo IAM Storage Object Viewer (roles/storage.objectViewer) al service account Compute Engine predefinito.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --role roles/storage.objectViewer

Esegui il deployment del controller JobSet

Il controller JobSet è responsabile della gestione dei job batch che eseguono l'addestramento del modello su GKE e la sua allocazione delle risorse viene modificata per gestire il carico di lavoro in modo efficiente. Assicurati che il launcher del job di addestramento esegua il deployment e utilizzi JobSet.

Per aumentare la richiesta di memoria a 1 Gi, il limite di memoria a 2 Gi e la richiesta di CPU a 1 per il container del gestore nella distribuzione JobSet, esegui il seguente comando patch:

kubectl patch -n jobset-system deploy jobset-controller-manager --type json \
    --patch '[{"op": "add", "path": "/spec/template/spec/containers/0/resources", "value": {"limits": {"memory": "2Gi"}, "requests": {"cpu": "1", "memory": "1Gi"}}}]'

Inizializza il driver CSI Multi-Tier Checkpointing

Questa sezione descrive come inizializzare il driver CSI Multi-Tier Checkpointing sui nodi in cui verranno eseguiti i tuoi workload. Il driver CSI è responsabile della gestione dell'archiviazione e della gestione dei checkpoint durante il processo di addestramento del modello.

Crea un CheckpointConfiguration

CheckpointConfiguration è una risorsa personalizzata Kubernetes che specifica le proprietà per il deployment del driver CSI Multi-Tier Checkpointing. Questa risorsa ha ambito a livello di cluster.

  1. Crea il seguente manifest checkpoint.yaml.

    apiVersion: checkpointing.gke.io/v1
kind: CheckpointConfiguration
metadata:
  name: MTC_CONFIG_NAME-configuration
spec:
    cloudStorageBucketName: GCS_BUCKET
    nodeSelector:
        node.kubernetes.io/instance-type: MACHINE_TYPE
    tolerations:
    - key: TOLERATION_KEY
        operator: Exists
        effect: NoSchedule
    inMemoryVolumeSize: IN_MEMORY_VOLUME_SIZE
    gcsFuseMountOptions:
    - implicit-dirs
    - metadata-cache:negative-ttl-secs:0
    - metadata-cache:ttl-secs:-1
    - metadata-cache:stat-cache-max-size-mb:-1
    - metadata-cache:type-cache-max-size-mb:-1
    - file-cache:max-size-mb:-1
    - file-cache:cache-file-for-range-read:true
    - file-system:kernel-list-cache-ttl-secs:0
    - file-cache:enable-parallel-downloads:true
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true

    Sostituisci quanto segue:

    • MTC_CONFIG_NAME: il nome di CheckpointConfiguration. Questo nome è globale per il cluster e non è specifico per il job.
    • GCS_BUCKET: il nome del bucket Cloud Storage in cui memorizzerai i dati del checkpoint. Utilizza il bucket che hai configurato nel passaggio Configura un bucket Cloud Storage con le autorizzazioni.

    • MACHINE_TYPE: il tipo di macchina per gli acceleratori corrispondenti. Il valore può essere uno dei seguenti:

      Per saperne di più sull'esecuzione di workload distribuiti su GPU con GKE, consulta Esecuzione di GPU multi-istanza. Per le TPU, consulta Crea il pool di nodi della sezione TPU.

    • TOLERATION_KEY: questo campo consente di programmare il driver CSI sui nodi con incompatibilità corrispondenti. Per ulteriori informazioni sul funzionamento dei taint sui diversi tipi di acceleratori, consulta queste pagine:

    • IN_MEMORY_VOLUME_SIZE: la dimensione della cache di checkpoint in memoria. Specifica la quantità e l'unità (ad esempio 200 Gi).Questo valore deve essere:

      • La dimensione del checkpoint locale per le TPU moltiplicata per 2,2
      • La dimensione del checkpoint locale per le GPU con un singolo peer moltiplicata per 6,6.

  2. Applica il manifest:

    kubectl apply -f checkpoint.yaml

  3. Controlla che il driver CSI sia in esecuzione:

    kubectl get pod -n gke-managed-checkpointing

    L'output dovrebbe essere simile al seguente. Ci saranno più voci, una per ogni nodo accelerato.

    NAME                                                          READY   STATUS    RESTARTS   AGE
multitier-driver-e2b033a7-a4e7-496a-87a3-ffd7fcc2e57b-2d4fz   5/5     Running   0          114s

Disinstalla il driver CSI Multi-Tier Checkpointing

Se vuoi annullare il deployment del driver CSI Multi-Tier Checkpointing, elimina la risorsa CheckpointConfiguration. Il controller Multi-Tier Checkpointing rimuove il driver CSI dai nodi. In questo modo vengono rimossi i dischi RAM e viene liberata memoria per altri carichi di lavoro. Ad esempio:

kubectl delete -f checkpoint.yaml

Gestire la conservazione dei dati e la garbage collection per i backup di Cloud Storage

Sei responsabile dell'implementazione dei criteri di conservazione per i backup dei checkpoint di Cloud Storage. Il checkpoint multilivello scrive i backup dei checkpoint solo in Cloud Storage e non li modifica o elimina mai.

Molti strumenti open source possono gestire la conservazione e la garbage collection, tra cui:

L'esempio seguente utilizza backup-warden in cui la directory backup è montata in una posizione di backup che utilizza Cloud Storage FUSE:

# Add --delete option to actually delete the backups, as is it only shows what would be deleted (dry-run)
backup-warden -p backup \
    --hourly 24 \
    --daily 7 \
    --weekly 5 \
    --monthly always \
    --yearly always \
    --prefer-recent

Aggiorna il manifest JobSet del workload

Aggiorna il manifest JobSet per il tuo job in modo da includere il volume del checkpoint su larga scala. I dettagli dipendono dal tuo workload.

Ad esempio, per estendere il JobSet di esempio da Esegui il deployment di TPU multislices in GKE, segui questi passaggi:

  1. Aggiungi le seguenti righe al contenitore jax-tpu.

    volumeMounts:
- name: checkpoint
  mountPath: CHECKPOINT_DIR

    Sostituisci CHECKPOINT_DIR con il percorso della directory dei checkpoint. Questa è la posizione in cui viene generato replicator.yaml e in cui Multi-Tier Checkpointing esegue l'operazione di salvataggio del checkpoint. Per maggiori informazioni, vedi Integrare il checkpointing multilivello nell'applicazione.

  2. Aggiungi le seguenti righe al campo spec.template.spec della specifica del job.

    volumes:
- name: checkpoint
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Integrare il checkpointing multilivello nell'applicazione

Per condividere informazioni sulle posizioni dei checkpoint e sulla preparazione della replica, modifica l'applicazione in modo che utilizzi il seguente protocollo per comunicare con il checkpoint multitier.

Avvio

Questa sezione descrive i passaggi iniziali necessari all'applicazione per interagire con il checkpointing multilivello.

Replicator è un componente principale del checkpointing multilivello, in esecuzione su ogni nodo come parte del driver CSI. Replicator gestisce la replica dei checkpoint nei livelli di archiviazione, dal disco RAM locale ai nodi peer e all'archiviazione esterna come Cloud Storage.

Il file replicator.yaml funge da piano di controllo dinamico tra il job di addestramento ML (codice del framework) e il componente Replicator. L'applicazione ML genera questo file a livello di programmazione sul volume locale (RAMDisk), accessibile sia dal job di addestramento sia dal servizio Replicator. Questo manifest consente al framework ML di fornire istruzioni di configurazione e gestione del ciclo di vita al replicatore, distinte dai parametri dell'infrastruttura statica (ad esempio, la frequenza di caricamento di Cloud Storage) definiti durante la configurazione del backend.

Per un esempio concreto di questa interazione, vedi:

La tua applicazione deve eseguire i seguenti passaggi durante l'avvio:

  1. Attendi finché il file replicator.yaml non è presente, il che indica che Replicator è pronto per essere configurato dalla tua applicazione. Il file replicator.yaml viene generato nella posizione CHECKPOINT_DIR che hai configurato nella sezione Aggiorna il manifest JobSet del carico di lavoro.

    Quando viene creato per la prima volta il job di addestramento del modello, il file replicator.yaml non esisterà e l'applicazione potrà procedere immediatamente. Tuttavia, se il job è stato riavviato (ad esempio, a causa di un errore o di un intervento manuale), il sistema potrebbe essere ancora in fase di elaborazione dell'istanza del job precedente e il replicator.yaml di questa istanza potrebbe essere ancora presente sul volume locale.

  2. L'applicazione o il job ML crea il file replicator.yaml con la configurazione simile alla seguente.

    Orbax

    job-name: orbax
framework: orbax
assume-data-parallelism: 3
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    PyTorch

    job-name: nemo
framework: pytorch.distributed
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    Questa configurazione di esempio ha i seguenti campi:

    • name: il nome del job di addestramento.
    • framework: il framework ML utilizzato dal job di addestramento.
    • node-rank: l'identificatore univoco del nodo corrente all'interno del job di addestramento distribuito. Rappresenta il rango del nodo che crea questo file. Ogni nodo che partecipa all'esecuzione avrà il proprio ranking.
    • nodes: il numero totale di nodi che partecipano al job di addestramento distribuito. Questo valore proviene dai metadati del pod. Anche il job di addestramento ML può visualizzare questo valore.
    • peer-ranks o peers-per-node: due modi alternativi per specificare la topologia di replica. Deve essere presente solo uno di questi due parametri.
      • peer-ranks: i ranghi espliciti dei nodi peer a cui devono essere replicati i dati del checkpoint del nodo corrente. In questo modo, puoi controllare con precisione quali nodi specifici fungono da partner di replica.
      • peers-per-node: il numero di nodi peer per nodo che il replicatore deve selezionare automaticamente per la replica.
    • backup-interval-minutes: la frequenza, in minuti, con cui i checkpoint vengono sottoposti a backup in Cloud Storage. Ti consigliamo di impostare questo valore su 30 minuti o più.

  3. Attendi che il nuovo file replicator.yaml venga eliminato dal sistema. Ciò indica che il replicatore è stato riavviato ed è stata eseguita la pulizia. Questo passaggio ti consente di evitare file obsoleti o temporanei sul volume locale quando l'applicazione esegue i passaggi nella sezione successiva.

Ripristinare dal checkpoint Ultimo stato noto

  1. Dopo l'inizializzazione del replicatore, il checkpoint multilivello crea un collegamento simbolico per ogni worker TPU o GPU. Questi link simbolici vengono creati nello stesso volume locale montato del file replicator.yaml, in cui il job salva i checkpoint.

    I link simbolici hanno il formato <job-name>-s{step}-n<node-rank>-w<worker-index>.restore.

  2. Ripristina ogni worker dal file .restore corrispondente. Per un esempio, consulta l'esempio di gestore dei checkpoint replicati di Orbax nella sezione successiva.

Salva il checkpoint

La tua applicazione esegue questi passaggi più volte durante l'avanzamento del job di addestramento. Le operazioni di salvataggio vengono eseguite nella posizione CHECKPOINT_DIR che hai configurato in Aggiorna il manifest JobSet del carico di lavoro.

Orbax

Crea checkpoint Orbax. La directory viene denominata con il numero del passaggio. Il replicatore rileva la directory del checkpoint appena creata, esegue la replica o il backup in base alle necessità ed esegue la pulizia automatica.

Per ulteriori informazioni su come utilizzare Orbax replicator checkpoint manager, consulta il file MaxtTest checkpointing. Per un esempio di interazione del servizio di replica, consulta il file MaxText max_utils.

PyTorch

Utilizza InClusterLocalCheckpointIO come pytorch_lightning.CheckpointIO personalizzato per attivare il checkpointing distribuito corretto con l'archiviazione locale. Il seguente comando di esempio abilita il checkpointing multilivello utilizzando un'implementazione di riferimento basata sul framework NVIDIA NeMo:

torchrun train.py <other_train_flags> \
    --local-ckpt-dir=CHECKPOINT_DIR \
    --local-ckpt-interval=20 \
    --job-name=JOB_NAME \
    --enable-high-scale-ckpt

Sostituisci quanto segue:

  • CHECKPOINT_DIR: il percorso della directory dei checkpoint.
  • JOB_NAME: il nome del workload del job di addestramento.

Upgrade dei cluster

Per gli upgrade del cluster, puoi eliminare e ricreare l'oggetto CheckpointConfiguration prima o dopo l'upgrade. Questa azione è necessaria perché il daemonset del driver di checkpointing dei nodi, implementato dinamicamente da questo oggetto, non verrà aggiornato automaticamente.

Se vuoi mantenere la stessa specifica del daemonset, non devi fare nulla.

Risoluzione dei problemi

Questa sezione fornisce indicazioni per la risoluzione dei problemi per risolvere i problemi relativi al checkpointing multilivello. Per la risoluzione dei problemi di archiviazione generali, consulta Risoluzione dei problemi di Cloud Storage in GKE.

Checkpointing multilivello non abilitato

Il seguente errore indica che il checkpoint multilivello non è abilitato sul tuo cluster:

error: unable to recognize "checkpoint.yaml": no matches for kind "CheckpointConfiguration" in version "checkpointing.gke.io/v1"

Potresti riscontrare questo errore dopo aver eseguito kubectl apply -f checkpoint.yaml nel passaggio Crea un CheckpointConfiguration.

Per risolvere il problema, controlla se hai abilitato il checkpoint multilivello sul cluster con il seguente comando:

gcloud container clusters describe CLUSTER_NAME \
    --project PROJECT_ID
    --location CLUSTER_LOCATION

Se il checkpoint a più livelli è abilitato, l'output dovrebbe essere simile al seguente:

addonsConfig:
  gcePersistentDiskCsiDriverConfig:
    enabled: true
  gcsFuseCsiDriverConfig:
    enabled: true
  highScaleCheckpointingConfig:
    enabled: true
  kubernetesDashboard:
    disabled: true
  networkPolicyConfig:
    disabled: true

Se il checkpointing multilivello non è attivato, aggiorna il cluster per attivarlo.

Il driver CSI Multi-Tier Checkpointing non è in grado di montare i volumi

Potresti riscontrare questo problema se il driver CSI non è in grado di montare il volume Cloud Storage. Potrebbero esserci più righe simili a questa.

kubectl get pod -n gke-managed-checkpointing
NAME                                                          READY   STATUS     RESTARTS   AGE
multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9   0/5     Init:0/1   0          6m32s

Per risolvere il problema, controlla gli eventi del pod del driver CSI, come mostrato nell'esempio seguente:

kubectl describe pod multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 -n gke-managed-checkpointing

Events:
  Type     Reason       Age                 From               Message
  ----     ------       ----                ----               -------
  Normal   Scheduled    17m                 default-scheduler  Successfully assigned gke-managed-checkpointing/multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 to gke-my-cluster-default-pool-353c773f-6d8q
  Warning  FailedMount  82s (x16 over 17m)  kubelet            MountVolume.SetUp failed for volume "gcs" : rpc error: code = PermissionDenied desc = failed to get GCS bucket "checkpointing-test-bucket": googleapi: Error 403: Caller does not have storage.objects.list access to the Google Cloud Storage bucket. Permission 'storage.objects.list' denied on resource (or it may not exist)., forbidden

Se il problema si verifica a causa dell'errore PermissionDenied del bucket Cloud Storage, come mostrato nell'esempio, puoi risolverlo configurando correttamente le autorizzazioni.

Passaggi successivi