Profilare il modello sulle VM Cloud TPU

La profilazione è uno dei principali strumenti per ottimizzare le prestazioni dei tuoi modelli su Cloud TPU. Lo strumento di profilazione principale si chiama XProf, disponibile nel repository GitHub OpenXLA/XProf. XProf supporta la profilazione di tutti i framework basati su XLA, tra cui JAX, Pytorch XLA e Tensorflow/Keras.

Acquisizione dei profili

Per utilizzare XProf, devi prima acquisire un profilo del modello. Esistono due modi per acquisire i profili:

  1. Acquisizione programmatica
  2. Acquisizione on demand (nota anche come acquisizione manuale)

Con l'acquisizione programmatica, devi annotare il codice del modello per specificare in quale punto del codice vuoi acquisire i profili. In genere, acquisisci un profilo per alcuni passaggi di addestramento o profili un blocco di codice specifico all'interno del modello. Esistono diversi modi per acquisire le tracce nei diversi framework JAX, Pytorch XLA e TensorFlow: avvio o interruzione della traccia basati su API o basati su context manager. Nei framework di livello superiore come MaxText, puoi attivare l'acquisizione dei profili semplicemente attivando un flag profiler=xplane durante l'esecuzione di MaxText.

Puoi utilizzare l'acquisizione dei profili on demand quando vuoi acquisire i profili in modo ad hoc o quando non hai attivato l'acquisizione programmatica dei profili. Questo può essere utile quando si verifica un problema con le metriche del modello durante l'esecuzione e si vogliono acquisire profili in quel momento per un determinato periodo di tempo per diagnosticare il problema.

Per attivare l'acquisizione di profili on demand, devi comunque avviare il server XProf all'interno del codice. L'attivazione di profiler.start_server avvierà un server XProf sul tuo workload ML che ascolta il trigger di acquisizione on demand per avviare l'acquisizione dei profili. Puoi attivare l'acquisizione di profili on demand tramite l'interfaccia utente di TensorBoard o tramite la CLI con lo strumento XProfiler.

Puoi scoprire come attivare l'acquisizione programmatica e on demand per diversi framework all'indirizzo:

Per l'acquisizione di profili sia programmatica sia on demand, devi specificare la posizione in cui vuoi archiviare i profili acquisiti. Puoi scegliere di archiviare i profili in una directory della tua VM TPU con un percorso simile a /profiles/run-name o utilizzare Cloud Storage con un percorso simile a gs://bucket-name/run-name/.

Con i percorsi menzionati, i tuoi profili verranno acquisiti in:

/profiles/run-name/plugins/profile/session1/

o

gs://bucket-name/run-name/plugins/profile/session1/.

Una sessione rappresenta un'acquisizione dei profili durante l'esecuzione del modello. Ad esempio, se acquisisci profili in una corsa di allenamento dai passaggi 1-3 e in seguito acquisisci profili dai passaggi 8-10, questi sono profili per la stessa corsa, ma la prima acquisizione dai passaggi 1-3 sarà session1 e la seconda acquisizione dai passaggi 8-10 sarà session2.

Le diverse sessioni saranno indicate con timestamp diversi in ogni esecuzione. Puoi acquisire profili in sessioni diverse in modo programmatico o on demand oppure una combinazione di entrambi. Questa funzionalità è utile quando vuoi confrontare i profili di diverse parti dell'addestramento del modello. Ad esempio, confrontando i profili dall'inizio dell'addestramento con quelli della fine dell'addestramento.

Visualizzazione dei profili con TensorBoard

Per visualizzare i profili XProf in TensorBoard, devi installare il plug-in tensorboard-plugin-profile.

Per visualizzare i profili acquisiti su Google Cloud, ti consigliamo di utilizzare TensorBoard su una VM TPU o lo strumento XProfiler per ospitare TensorBoard su una VM o un pod GKE.

Se hai acquisito i profili sulla VM TPU, puoi visualizzarli utilizzando TensorBoard su quella VM TPU.

Per installare XProf e TensorBoard su una VM TPU, puoi eseguire questi comandi:

pip install tensorboard_plugin_profile tensorboard
pip install xprof

Per visualizzare i profili acquisiti con TensorBoard, esegui:

$ tensorboard --logdir=profiles/run-name

o

$ tensorboard --logdir=profiles

Per caricare più profili da più esecuzioni, indirizza TensorBoard alla directory radice che contiene tutte le esecuzioni e i profili anziché alla sottodirectory che contiene i profili di un'esecuzione specifica.

Libreria Cloud-Diagnostics-XProf per visualizzare i profili TensorBoard su Google Cloud

cloud-diagnostics-xprof (nota anche come libreria XProfiler) semplifica l'hosting di TensorBoard e la visualizzazione dei profili su Google Cloud. Per saperne di più, consulta il repository GitHub cloud-diagnostics-xprof.

Il repository cloud-diagnostics-xprof offre i seguenti miglioramenti rispetto all'esecuzione locale di XProf e TensorBoard:

  • Configurazione e creazione di pacchetti delle dipendenze di XProf e TensorBoard
  • Archivia i tuoi profili in Cloud Storage, il che può essere utile per la conservazione a lungo termine e l'analisi post-esecuzione (i profili locali acquisiti verranno eliminati al termine dell'esecuzione da parte del ricercatore)
  • Caricamento rapido di profili di grandi dimensioni e di più profili mediante il provisioning di TensorBoard su una VM Compute Engine o un pod GKE, con la possibilità di modificare il tipo di macchina in base alle esigenze dell'utente per velocità di caricamento e costi
  • Crea un link per la condivisione dei profili e la collaborazione con i membri del team e gli ingegneri di Google
  • Profilazione on demand più semplice dei carichi di lavoro su GKE e Compute Engine per scegliere qualsiasi host che esegue il carico di lavoro per acquisire i profili

Prima di utilizzare la libreria XProfiler, devi acquisire i profili per il codice del tuo carico di lavoro in modo programmatico o avviando il server di profilazione, in modo da poter eseguire l'acquisizione dei profili on demand in un secondo momento.

Per configurare la libreria XProfiler, devi configurare Google Cloud CLI e un ambiente virtuale Python, quindi devi solo eseguire:

pip install cloud-diagnostics-xprof

Verranno installate tutte le dipendenze necessarie di XProf e TensorBoard.

Successivamente, per creare una VM o un pod GKE per ospitare TensorBoard, esegui questo comando:

xprofiler create -z $ZONE -l $GCS_PATH

o

xprofiler create --GKE -z $ZONE -l $GCS_PATH

Sostituisci $ZONE con una zona qualsiasi e $GCS_PATH con il percorso delle tracce del profilo. Puoi specificare la directory principale che contiene tutte le tracce del profilo per più esecuzioni o un insieme specifico di tracce del profilo per un'esecuzione.

Ad esempio, se acquisisci profili nelle seguenti directory:

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

Puoi impostare il percorso GCS sulla directory principale (gs://bucket-name):

xprofiler create -z $ZONE -l gs://bucket-name/

Nell'interfaccia utente di TensorBoard, vedrai tutti i profili per tutte le esecuzioni e le sessioni come run1/session1, run1/session2 e run2/session1.

Per impostazione predefinita, xprofiler create crea una VM Compute Engine, in particolare una VM c4-highmem-8. Puoi modificare il tipo di macchina con il flag -m. Se vuoi creare un'istanza TensorBoard su un pod GKE anziché su una VM di Compute Engine, puoi trasmettere il flag –GKE a xprofiler create. L'hosting di un'istanza di TensorBoard su un pod GKE può semplificare la gestione di TensorBoard insieme al resto del carico di lavoro di cui è stato eseguito il deployment su GKE.

La VM o il pod GKE che ospita TensorBoard rende il caricamento di profili di grandi dimensioni e di più profili molto più veloce su Google Cloud rispetto a TensorBoard ospitato localmente. In base al nostro benchmarking, vedrai profili dell'ordine di 1 GB di carico entro pochi minuti per il primo caricamento utilizzando la VM c4-highmem-8 predefinita. Inoltre, puoi scegliere il tipo di macchina giusto in base alle tue esigenze di prestazioni e costi.

Dopo aver eseguito xprofiler create, vedrai un output simile al seguente:

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

La prima opzione è un link su cui puoi fare clic per visualizzare i tuoi profili XProf su TensorBoard. Questo link consente la condivisione dei profili con il tuo team e anche con gli ingegneri di Google che si occupano dell'ottimizzazione delle prestazioni su Google Cloud.

Sei tu a controllare chi ha accesso al link in base alle autorizzazioni impostate per il bucket Cloud Storage in cui sono archiviati i dati del profilo. Se per qualche motivo il link non funziona, forniamo anche un modo per connettersi all'istanza TensorBoard utilizzando SSH per visualizzare i profili utilizzando il comando xprofiler connect.

Se hai abilitato il server Cloud Profiler nel codice del tuo workload (operazione richiesta sia per le acquisizioni programmatiche che per quelle on demand) e vuoi eseguire la profilazione on demand, puoi farlo in due modi:

a. Fai clic sul pulsante Acquisisci profilo su TensorBoard. Ora puoi scegliere qualsiasi host del dispositivo su cui viene eseguito il tuo carico di lavoro per la profilazione. Supportiamo l'acquisizione on demand per i carichi di lavoro in esecuzione su GKE o Compute Engine.

b. Utilizza il comando di acquisizione XProfiler, fornendo informazioni come zona, bucket Cloud Storage, framework, nomi di VM/pod host e durata dell'acquisizione in millisecondi. Si tratta delle stesse informazioni che devi inserire nell'interfaccia utente di TensorBoard.

Per maggiori dettagli su come utilizzare la libreria cloud-diagnostics-xprof, consulta questa pagina GitHub.