Aumentare la velocità del checkpoint e ridurre i costi con Nebula

Articolo
09/01/2024

Informazioni su come aumentare la velocità dei checkpoint e ridurre i costi dei checkpoint per i modelli di training di Azure Machine Learning di grandi dimensioni usando Nebula.

Panoramica

Nebula è uno strumento di checkpoint rapido, semplice, senza disco e compatibile con il modello in Azure Container per PyTorch (ACPT). Nebula offre una soluzione di checkpoint semplice e ad alta velocità per i processi di training di modelli su larga scala distribuiti tramite PyTorch. Usando le tecnologie di elaborazione distribuite più recenti, Nebula può ridurre i tempi di checkpoint da ore a secondi, risparmiando potenzialmente il 95% al 99,9% del tempo. I processi di training su larga scala possono trarre notevole vantaggio dalle prestazioni di Nebula.

Per rendere Nebula disponibile per i processi di training, importare il nebulaml pacchetto Python nello script. Nebula ha piena compatibilità con diverse strategie di training di PyTorch distribuite, tra cui PyTorch Lightning, DeepSpeed e altro ancora. L'API Nebula offre un modo semplice per monitorare e visualizzare i cicli di vita dei checkpoint. Le API supportano vari tipi di modello e garantiscono la coerenza e l'affidabilità dei checkpoint.

Importante

Il nebulaml pacchetto non è disponibile nell'indice del pacchetto Python PyPI pubblico. È disponibile solo nell'ambiente curato azure Container for PyTorch (ACPT) in Azure Machine Learning. Per evitare problemi, non tentare di eseguire l'installazione nebulaml da PyPI o usando il pip comando .

In questo documento si apprenderà come usare Nebula con ACPT in Azure Machine Learning per registrare rapidamente i processi di training del modello. Si apprenderà anche come visualizzare e gestire i dati del checkpoint Nebula. Si apprenderà anche come riprendere i processi di training del modello dall'ultimo checkpoint disponibile in caso di interruzione, errore o chiusura di Azure Machine Learning.

Perché l'ottimizzazione dei checkpoint per il training di modelli di grandi dimensioni è importante

Man mano che i volumi di dati aumentano e i formati di dati diventano più complessi, anche i modelli di Machine Learning sono diventati più sofisticati. Il training di questi modelli complessi può risultare complesso a causa dei limiti di capacità della memoria GPU e dei tempi di training lunghi. Di conseguenza, il training distribuito viene spesso usato quando si usano set di dati di grandi dimensioni e modelli complessi. Tuttavia, le architetture distribuite possono riscontrare errori imprevisti e errori dei nodi, che possono diventare sempre più problematici man mano che aumenta il numero di nodi in un modello di Machine Learning.

I checkpoint consentono di attenuare questi problemi salvando periodicamente uno snapshot dello stato completo del modello in un determinato momento. In caso di errore, questo snapshot può essere usato per ricompilare il modello allo stato al momento dello snapshot in modo che il training possa riprendere da quel punto.

Quando si verificano errori o terminazioni delle operazioni di training del modello di grandi dimensioni, i data scientist e i ricercatori possono ripristinare il processo di training da un checkpoint salvato in precedenza. Tuttavia, tutti gli progressi compiuti tra il checkpoint e la terminazione vengono persi perché i calcoli devono essere riesecuti per recuperare i risultati intermedi non salvati. Gli intervalli di checkpoint più brevi potrebbero contribuire a ridurre questa perdita. Il diagramma illustra il tempo sprecato tra il processo di training dai checkpoint e la terminazione:

Tuttavia, il processo di salvataggio dei checkpoint stesso può generare un sovraccarico significativo. Il salvataggio di un checkpoint di dimensioni TB può spesso diventare un collo di bottiglia nel processo di training, con il processo di checkpoint sincronizzato che blocca il training per ore. In media, i sovraccarichi correlati al checkpoint possono tenere conto del 12% del tempo totale di training e possono aumentare fino al 43% (Maeng et al., 2021).

Per riepilogare, la gestione dei checkpoint del modello di grandi dimensioni comporta un elevato carico di archiviazione e sovraccarichi del tempo di ripristino dei processi. I checkpoint frequenti vengono salvati, combinati con le riprese dei processi di training dai checkpoint disponibili più recenti, diventano una grande sfida.

Nebulosa al salvataggio

Per eseguire il training efficace di modelli distribuiti di grandi dimensioni, è importante avere un modo affidabile ed efficiente per risparmiare e riprendere lo stato di training, riducendo al minimo la perdita di dati e lo spreco di risorse. Nebula consente di ridurre i tempi di salvataggio dei checkpoint e le richieste di ore gpu per i processi di training di Azure Machine Learning di grandi dimensioni, offrendo una gestione dei checkpoint più rapida e semplice.

Con Nebula è possibile:

Aumentare la velocità dei checkpoint fino a 1000 volte con un'API semplice che funziona in modo asincrono con il processo di training. Nebula può ridurre i tempi di checkpoint da ore a secondi- una potenziale riduzione del 95% al 99%.

Questo esempio mostra la riduzione del checkpoint e del tempo di training end-to-end per quattro checkpoint che salvano i processi di training Hugging Face GPT2, GPT2-Large e GPT-XL. Per il checkpoint Hugging Face GPT2-XL di medie dimensioni salva (20,6 GB), Nebula ha raggiunto una riduzione del tempo del 96,9% per un checkpoint.

L'aumento della velocità del checkpoint può comunque aumentare con le dimensioni del modello e i numeri della GPU. Ad esempio, il test di un checkpoint del punto di training di 97 GB su GPU Nvidia da 128 A100 può ridursi da 20 minuti a 1 secondo.
Ridurre i costi di training e calcolo end-to-end per i modelli di grandi dimensioni riducendo al minimo il sovraccarico del checkpoint e riducendo il numero di ore della GPU sprecate per il ripristino del processo. Nebula salva i checkpoint in modo asincrono e sblocca il processo di training per ridurre il tempo di training end-to-end. Consente anche di salvare i checkpoint più frequenti. In questo modo, è possibile riprendere il training dal checkpoint più recente dopo qualsiasi interruzione e risparmiare tempo e denaro sprecato per il ripristino del processo e le ore di training gpu.
Garantire la compatibilità completa con PyTorch. Nebula offre una compatibilità completa con PyTorch e offre un'integrazione completa con framework di training distribuiti, tra cui DeepSpeed (>=0.7.3) e PyTorch Lightning (>=1.5.0). È anche possibile usarlo con destinazioni di calcolo di Azure Machine Learning diverse, ad esempio Calcolo di Azure Machine Learning o servizio Azure Kubernetes.
Gestire facilmente i checkpoint con un pacchetto Python che consente di elencare, ottenere, salvare e caricare i checkpoint. Per mostrare il ciclo di vita del checkpoint, Nebula fornisce anche log completi su studio di Azure Machine Learning. È possibile scegliere di salvare i checkpoint in un percorso di archiviazione locale o remoto
- Archiviazione BLOB di Azure
- Azure Data Lake Storage
- NFS
e accedervi in qualsiasi momento con poche righe di codice.

Prerequisiti

Una sottoscrizione di Azure e un'area di lavoro di Azure Machine Learning. Per altre informazioni sulla creazione delle risorse dell'area di lavoro, vedere Creare risorse dell'area di lavoro
Destinazione di calcolo di Azure Machine Learning. Per altre informazioni sulla creazione della destinazione di calcolo, vedere Gestire il training e distribuire i calcoli
Script di training che usa PyTorch.
Ambiente ACPT curato (Azure Container per PyTorch). Vedere Ambienti curati per ottenere l'immagine ACPT. Informazioni su come usare l'ambiente curato

Come usare Nebula

Nebula offre un'esperienza di checkpoint veloce e semplice, direttamente nello script di training esistente. I passaggi per avviare rapidamente Nebula includono:

Uso dell'ambiente ACPT
Inizializzazione di Nebula
Chiamata di API per salvare e caricare checkpoint

Uso dell'ambiente ACPT

Azure Container per PyTorch (ACPT), un ambiente curato per il training del modello PyTorch, include Nebula come pacchetto Python predefinito e dipendente. Per altre informazioni sull'immagine ACPT, vedere Azure Container for PyTorch (ACPT) per visualizzare l'ambiente curato e Abilitazione di Deep Learning con Azure Container per PyTorch in Azure Machine Learning .

Inizializzazione di Nebula

Per abilitare Nebula con l'ambiente ACPT, è sufficiente modificare lo script di training per importare il nebulaml pacchetto e quindi chiamare le API Nebula nelle posizioni appropriate. È possibile evitare la modifica dell'SDK o dell'interfaccia della riga di comando di Azure Machine Learning. È anche possibile evitare di modificare altri passaggi per eseguire il training del modello di grandi dimensioni nella piattaforma Azure Machine Learning.

Nebula richiede l'inizializzazione per l'esecuzione nello script di training. Nella fase di inizializzazione specificare le variabili che determinano la posizione e la frequenza di salvataggio del checkpoint, come illustrato in questo frammento di codice:

  import nebulaml as nm
  nm.init(persistent_storage_path=<YOUR STORAGE PATH>) # initialize Nebula

Nebula è stata integrata in DeepSpeed e PyTorch Lightning. Di conseguenza, l'inizializzazione diventa semplice e semplice. Questi esempi illustrano come integrare Nebula negli script di training.

Importante

Il salvataggio dei checkpoint con Nebula richiede memoria per archiviare i checkpoint. Assicurarsi che la memoria sia maggiore di almeno tre copie dei checkpoint.

Se la memoria non è sufficiente per contenere checkpoint, è consigliabile configurare una variabile NEBULA_MEMORY_BUFFER_SIZE di ambiente nel comando per limitare l'uso della memoria per ogni nodo durante il salvataggio dei checkpoint. Quando si imposta questa variabile, Nebula userà questa memoria come buffer per salvare i checkpoint. Se l'utilizzo della memoria non è limitato, Nebula userà la memoria il più possibile per archiviare i checkpoint.

Se più processi vengono eseguiti nello stesso nodo, la memoria massima per il salvataggio dei checkpoint sarà la metà del limite diviso per il numero di processi. Nebula userà l'altra metà per il coordinamento multi-processo. Ad esempio, se si vuole limitare l'utilizzo della memoria per ogni nodo a 200 MB, è possibile impostare la variabile di ambiente come export NEBULA_MEMORY_BUFFER_SIZE=200000000 (in byte, circa 200 MB) nel comando. In questo caso, Nebula userà solo 200 MB di memoria per archiviare i checkpoint in ogni nodo. Se sono presenti 4 processi in esecuzione nello stesso nodo, Nebula userà 25 MB di memoria per ogni processo per archiviare i checkpoint.

Chiamata di API per salvare e caricare checkpoint

Nebula fornisce API per gestire i salvataggi dei checkpoint. È possibile usare queste API negli script di training, analogamente all'API PyTorch torch.save() . Questi esempi illustrano come usare Nebula negli script di training.

Visualizzare le cronologie dei checkpoint

Al termine del processo di training, passare al riquadro Processo Name> Outputs + logs . Nel pannello sinistro espandere la cartella Nebula e selezionare questa opzione per visualizzare checkpointHistories.csv informazioni dettagliate sui salvataggi del checkpoint Nebula: durata, velocità effettiva e dimensioni del checkpoint.

Esempi

Questi esempi illustrano come usare Nebula con tipi di framework diversi. È possibile scegliere l'esempio più adatto allo script di training.

Per abilitare la compatibilità nebulosa completa con gli script di training basati su PyTorch, modificare lo script di training in base alle esigenze.

Prima di tutto, importare il pacchetto richiesto nebulaml :

  # Import the Nebula package for fast-checkpointing 
  import nebulaml as nm

Per inizializzare Nebula, chiamare la nm.init() funzione in main(), come illustrato di seguito:

  # Initialize Nebula with variables that helps Nebula to know where and how often to save your checkpoints
  persistent_storage_path="/tmp/test",
  nm.init(persistent_storage_path, persistent_time_interval=2)

Per salvare i checkpoint, sostituire l'istruzione originale torch.save() per salvare il checkpoint con Nebula. Assicurarsi che l'istanza del checkpoint inizi con "global_step", ad esempio "global_step500" o "global_step1000":
```
checkpoint = nm.Checkpoint('global_step500')  
checkpoint.save('<CKPT_NAME>', model)  
```
Nota

<'CKPT_TAG_NAME'> è l'ID univoco per il checkpoint. Un tag è in genere il numero di passaggi, il numero di periodo o qualsiasi nome definito dall'utente. Il parametro facoltativo facoltativo <'NUM_OF_FILES'> specifica il numero di stato salvato per questo tag.
Caricare il checkpoint valido più recente, come illustrato di seguito:
```
latest_ckpt = nm.get_latest_checkpoint()
p0 = latest_ckpt.load(<'CKPT_NAME'>)
```
Poiché un checkpoint o uno snapshot può contenere molti file, è possibile caricarne uno o più in base al nome. Con il checkpoint più recente, lo stato di training può essere ripristinato nello stato salvato dall'ultimo checkpoint.

Altre API possono gestire la gestione dei checkpoint
- elencare tutti i checkpoint
- ottenere i checkpoint più recenti
```
# Managing checkpoints
## List all checkpoints
ckpts = nm.list_checkpoints()
## Get Latest checkpoint path
latest_ckpt_path = nm.get_latest_checkpoint_path("checkpoint", persisted_storage_path)
```

Uno script di training basato su DeepSpeed (>=0.7.3) può usare Nebula, se si abilita Nebula nel ds_config.json file di configurazione, come illustrato di seguito:

"nebula": {
      "enabled": true,
      "persistent_storage_path": "<YOUR STORAGE PATH>",
      "persistent_time_interval": 100,
      "num_of_version_in_retention": 2,
      "enable_nebula_load": true
  }

Questa funzione di frammenti JSON funziona come la nebulaml.init() funzione .

L'inizializzazione con ds_config.json la configurazione dei file abilita Nebula, che consente di salvare i checkpoint a sua volta. Il metodo di salvataggio DeepSpeed originale, con l'API model_engine.save_checkpoint()di checkpoint del modello, usa automaticamente Nebula. Questo metodo di salvataggio evita la necessità di modifica del codice.

PyTorch Lightning (Nebula supporta automaticamente i checkpoint versione >=1.5.0) quando viene usato Trainer. Come spesso si salvano i checkpoint con comportamenti personalizzati per il controllo con granularità fine, PyTorch Lightning offre due modi per salvare il checkpoint: i salvataggi condizionali con ModelCheckpoint()e i salvataggi manuali con trainer.save_checkpoint(). Queste tecniche si applicano agli script di training di PyTorch (>=0.15.0).

Se si usa ModelCheckpoint per salvare in modo condizionale i checkpoint, è possibile usare NebulaCallback anziché ModelCheckpoint per l'inizializzazione.

# import Nebula package
import nebulaml as nm

# define NebulaCallback
config_params = dict()
config_params["persistent_storage_path"] = "<YOUR STORAGE PATH>"
config_params["persistent_time_interval"] = 10

nebula_checkpoint_callback = nm.NebulaCallback(
   ****, # Original ModelCheckpoint params
   config_params=config_params, # customize the config of init nebula
)

nm.NebulaCheckpointIO() Aggiungere quindi come plug-in a Trainere modificare i trainer.save_checkpoint() parametri di archiviazione come illustrato di seguito:

trainer = Trainer(plugins=[nm.NebulaCheckpointIO()],   # add NebulaCheckpointIO as a plugin
                  callbacks=[nebula_checkpoint_callback]) # use NebulaCallback as a plugin

Se si usa trainer.save_checkpoint() per salvare manualmente i checkpoint, è possibile usare il NebulaCheckpointIO plug-in in Trainere modificare i parametri di archiviazione in trainer.save_checkpoint() come indicato di seguito:

# import Nebula package
import nebulaml as nm

# initialize Nebula
nm.init(persistent_storage_path=<YOUR STORAGE PATH>) 

trainer = Trainer(plugins=[nm.NebulaCheckpointIO()])  # add NebulaCheckpointIO as a plugin

# Saving checkpoints
storage_options = {}
storage_options['is_best'] = True
storage_options['persist_path'] = "/tmp/tier3/checkpoint"

trainer.save_checkpoint("example.ckpt",
  storage_options=storage_options, # customize the config of Nebula saving checkpoint
)

Carica checkpoint

Carichiamo checkpoint coerenti con PyTorch e PyTorch Lightning. L'unica modifica specifica il percorso di archiviazione.

Per caricare il checkpoint più recente, MyLightningModule.load_from_checkpoint() funziona ancora, come illustrato di seguito:

  persistent_path = "/tmp/tier3/checkpoint"
  latest_ckpt_path = nebulaml.get_latest_checkpoint_path("checkpoint", persist_path)
  model = MyLightningModule.load_from_checkpoint(latest_ckpt_path)

Se è stato usato nebula_checkpoint_callback in Trainer(), lo script originale non modificato, come illustrato, funziona comunque:

  trainer = Trainer(
      default_root_dir=tmpdir, 
      max_steps=100,
      plugins=[NebulaCheckpointIO()],
      callbacks=[nebula_checkpoint_callback],
  )
  trainer.fit(model, ckpt_path="/path/example.ckpt")

Condividi tramite