Profila 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 ed è disponibile nel repository GitHub OpenXLA/XProf. XProf supporta la profilazione di tutti i framework basati su XLA, tra cui JAX, PyTorch e Tensorflow/Keras.
Acquisizione dei profili
Per utilizzare XProf, devi prima acquisire un profilo del tuo modello. Esistono due modi per acquisire i profili:
- Acquisizione programmatica
- Acquisizione on demand (nota anche come acquisizione manuale)
Con l'acquisizione programmatica, devi annotare il codice del modello per specificare dove nel 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 e TensorFlow, ovvero l'avvio o l'arresto della traccia basato su API o basato su gestore di contesto. In framework di livello superiore come MaxText, puoi attivare l'acquisizione del profilo 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 dei profili programmatica. Questo può essere utile quando riscontri un problema con le metriche del modello durante l'esecuzione e vuoi acquisire i profili in quel momento per un certo periodo di tempo per diagnosticare il problema.
Per attivare l'acquisizione del profilo on demand, devi comunque avviare il server XProf
all'interno del codice. L'abilitazione di profiler.start_server avvierà un server XProf sul tuo workload ML che ascolta il trigger di acquisizione on demand per iniziare ad acquisire i profili. Puoi attivare l'acquisizione del profilo on demand tramite l'interfaccia utente di Tensorboard o tramite l'interfaccia a riga di comando con lo strumento XProfiler.
Puoi scoprire come attivare l'acquisizione programmatica e on demand per diversi framework all'indirizzo:
Sia per l'acquisizione dei profili programmatica che 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 sopra 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 di profili durante un'esecuzione del modello. Ad esempio, se acquisisci i profili in un'esecuzione di addestramento dai passaggi 1-3 e in un secondo momento acquisisci i profili dai passaggi 8-10, questi sono i profili per la stessa esecuzione, ma la prima acquisizione dai passaggi 1-3 sarà session1 e la seconda acquisizione dai passaggi 8-10 sarà session2.
Le diverse sessioni verranno indicate con date diverse sotto ogni corsa. Puoi acquisire i profili in sessioni diverse in modo programmatico, on demand o con una combinazione di entrambi i metodi. Questa opzione è utile quando vuoi confrontare i profili di diverse parti dell'addestramento del modello. Ad esempio, confrontando i profili all'inizio dell'addestramento con quelli alla fine.
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 di utilizzare 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 sulla 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 principale che contiene tutte le esecuzioni e i profili, anziché alla sottodirectory che contiene i profili di un'esecuzione specifica.
Libreria Cloud-Diagnostics-XProf per la visualizzazione dei profili Tensorboard su Google Cloud
La 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 pacchettizzazione delle dipendenze di XProf e TensorBoard
- Archivia i tuoi profili in Cloud Storage, 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 eseguendo il provisioning di TensorBoard su una VM Compute Engine o un pod GKE, con l'opzione di modificare il tipo di macchina in base alle esigenze dell'utente in termini di 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 workload su GKE e Compute Engine per scegliere qualsiasi host che esegue il workload per acquisire i profili
Prima di utilizzare la libreria XProfiler, devi acquisire i profili per il codice del workload in modo programmatico o avviando il server del profilo, in modo da poter eseguire l'acquisizione del profilo on demand in un secondo momento.
Per configurare la libreria XProfiler, devi aver configurato Google Cloud CLI e un ambiente virtuale Python, quindi devi solo eseguire:
pip install cloud-diagnostics-xprof
In questo modo verranno installate tutte le dipendenze necessarie di XProf e TensorBoard.
Poi, per creare una VM o un pod GKE per ospitare TensorBoard, esegui il seguente comando:
xprofiler create -z $ZONE -l $GCS_PATH
o
xprofiler create --GKE -z $ZONE -l $GCS_PATH
Sostituisci $ZONE con qualsiasi zona e $GCS_PATH con il percorso delle tracce del profilo. Puoi specificare la directory root 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 i 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 di Cloud Storage sulla directory root (gs://bucket-name):
xprofiler create -z $ZONE -l gs://bucket-name/
Nella UI di TensorBoard, vedrai tutti i profili per tutte le esecuzioni e le sessioni come run1/session1, run1/session2, run2/session1.
Per impostazione predefinita, xprofiler create creerà una VM di Compute Engine,
nello specifico 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 passare il flag –GKE a xprofiler create. L'hosting di un'istanza TensorBoard su un pod GKE può semplificare la gestione di TensorBoard insieme al resto del workload di cui è stato eseguito il deployment su GKE.
La VM o il pod GKE che ospita TensorBoard velocizza notevolmente il caricamento di profili di grandi dimensioni e di più profili su Google Cloud rispetto a TensorBoard ospitato localmente. In base al nostro benchmarking, vedrai profili con un carico dell'ordine di 1 GB 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 una TensorBoard. Questo link consente di condividere i profili con il tuo team e con gli ingegneri Google che aiutano a ottimizzare le 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 tuo profilo. Se per qualche motivo il link non funziona, forniamo anche un modo per connetterti all'istanza TensorBoard utilizzando SSH per visualizzare i tuoi profili utilizzando il comando xprofiler connect.
Se hai attivato il server Cloud Profiler nel codice del workload (obbligatorio per le acquisizioni programmatiche e 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 di dispositivi su cui è in esecuzione il tuo workload per la profilazione. Supportiamo l'acquisizione on demand per i workload in esecuzione su GKE o Compute Engine.
b. Utilizza il comando di acquisizione di 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.