Risoluzione dei problemi di TensorFlow - TPU

Questa guida, insieme alle Domande frequenti, fornisce una guida alla risoluzione dei problemi per gli utenti che stanno addestrando modelli TensorFlow su Cloud TPU. Se stai risolvendo problemi di addestramento Pytorch o JAX, puoi consultare i documenti per la risoluzione dei problemi relativi a tali framework:

Per guide più generali su come utilizzare Cloud TPU, vedi:

Panoramica

I problemi comuni riscontrati con le Cloud TPU rientrano nelle seguenti categorie:

  1. Problemi di connessione alla TPU

  2. Debug degli errori comuni

  3. Ridurre l'utilizzo della memoria

  4. Migliorare la velocità di addestramento

  5. Debug dei cali di accuratezza del modello

Errore di connessione al server TPU

Questa sezione descrive come risolvere i problemi relativi a casi in cui TensorFlow smette di rispondere o stampa un errore durante la connessione alla TPU. Il passaggio della compilazione del grafico TPU può richiedere molto tempo per i modelli di grandi dimensioni, quindi lascia che lo script venga eseguito per almeno cinque minuti prima di concludere che non risponde più.

Il primo passaggio consiste nel verificare se il problema riguarda il server stesso o la pipeline di addestramento TensorFlow. Per farlo, esegui il tutorial MNIST utilizzando l'URL del tuo server TPU e verifica che funzioni correttamente. Se i problemi di connessione con il tutorial MNIST persistono, questo conferma che si tratta di un problema con il server TPU. In questo caso:

  1. Esegui il comando seguente per elencare le TPU disponibili. Sostituzione di zone e project-id con l'ID progetto e zona.

    (vm)$ gcloud compute tpus list --zone zone --project project-id

    L'output viene stampato, ad esempio: 

    NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

  2. Verifica di trasferire il valore corretto a --tpu (demo-tpu nell'esempio riportato sopra) e che questa TPU sia indicata come READY.

  3. Se la tua TPU non è elencata come READY o continui a riscontrare problemi di connessione, riavvia manualmente il server con:

    (vm)$ gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME

    Nell'esempio precedente, $TPU_SERVER_NAME è demo-tpu. Potrebbero essere necessari diversi minuti.

  4. Esegui di nuovo il comando ... tpus list sopra e attendi che lo stato della TPU sia nello stato READY. L'operazione potrebbe richiedere alcuni minuti.

  5. Prova a eseguire di nuovo il tutorial MNIST.

  6. Se i problemi di esecuzione del tutorial MNIST persistono, chiedi aiuto utilizzando uno dei meccanismi descritti in Richiedere assistenza.

Se l'esempio MNIST viene eseguito correttamente, ma il modello continua a non rispondere, è probabile che il problema sia relativo alla pipeline di addestramento. Per eseguire il debug, inizia sostituendo la TPUStrategia nel codice con la strategia predefinita. Quando utilizzi la strategia predefinita, ovunque utilizzi strategy.scope() o strategy.run(), il modello viene eseguito su CPU (o GPU se presente) anziché su TPU. Se il modello viene eseguito su CPU e non su TPU, deve esserci un problema specifico della TPU. Se il problema persiste, ti consigliamo di eseguire il debug del problema sulla CPU.

Perdita di connessione ssh durante l'addestramento

La connessione ssh a Cloud TPU potrebbe scadere durante un addestramento a lunga esecuzione (in particolare se utilizzi Cloud Shell). A quel punto, non è presente alcun output sulla console TPU e potrebbe sembrare che la TPU abbia interrotto l'addestramento. Per evitare questo problema, esegui la sessione di addestramento con un multiplexer terminale o uno strumento di gestione delle sessioni come tmux o screen. Il collegamento ssh rimarrà attivo indipendentemente dalla durata dell'addestramento.

Eseguire il debug degli errori comuni

Impossibile creare una TPU

Quando crei una Cloud TPU, potresti riscontrare il seguente errore:

googleapiclient.errors.HttpError: < HttpError 403 when requesting https://content-tpu.googleapis.com/v1/projects/{PROJECT}/locations/{ZONE}/nodes/{TPU_NAME}?alt=json returned "Request had insufficient authentication scopes."

Questo è un problema di autorizzazioni e può essere risolto eseguendo questo comando:

gcloud auth login --update-adc

Questo comando aggiorna le Credenziali predefinite dell'applicazione (ADC) e dovrebbe risolvere il problema. Per saperne di più, vedi gcloud auth login.

Impossibile utilizzare il file system locale

Messaggio di errore

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

Framework e configurazioni interessati

Questo messaggio può verificarsi durante l'addestramento con TensorFlow utilizzando l'architettura dei nodi TPU.

Dettagli

Tutti i file di input e la directory dei modelli devono utilizzare un percorso del bucket di spazio di archiviazione sul cloud (gs://bucket-name/...) e questo bucket deve essere accessibile dal server TPU. Tieni presente che tutte le elaborazioni dei dati e il checkpoint dei modelli vengono eseguite sul server TPU, non sulla macchina locale. Per informazioni su come configurare correttamente spazio di archiviazione sul cloud per l'utilizzo con TPU, consulta la guida Connessione a bucket Cloud Storage.

Tipo di dati non supportato

Messaggio di errore

TypeError: DataType is not a supported TPU infeed type.

Framework e configurazioni interessati

Questo messaggio può verificarsi durante l'addestramento con TensorFlow utilizzando l'architettura dei nodi TPU.

Dettagli

Attualmente, solo i tipi di dati tf.float32, tf.int32, tf.bfloat16 e tf.bool sono supportati sulla TPU. Altri tipi di dati comuni, come tf.uint8, tf.string e tf.int64, devono essere convertiti in uno dei tipi di dati supportati durante la pre-elaborazione dei dati (ovvero nella pipeline tf.data.Dataset).

Vedi un esempio di conversione nella funzione decode_image utilizzata nell'addestramento MNIST.

Forme dinamiche non supportate

Messaggio di errore

ValueError: shape [Shape] must have a fixed size for dimension
d that is known at graph construction time.

Framework e configurazioni interessati

Questo messaggio viene visualizzato solo durante la compilazione XLA con TensorFlow.

Dettagli

Per eseguire un modello sulla TPU, TensorFlow lo compila utilizzando il compilatore XLA. Sebbene questo passaggio di compilazione migliori notevolmente la velocità di addestramento e l'utilizzo della memoria, le forme (dimensioni delle dimensioni) di tutti i tensori nel grafico devono essere note al momento della compilazione del grafico. Se non è possibile determinare alcuna forma al momento della compilazione, la compilazione TPU ha esito negativo con un errore come quello riportato sopra.

Un'operazione comune che restituisce una forma dinamica è dataset.batch(batch_size), poiché il numero di campioni rimanenti in un flusso potrebbe essere inferiore alle dimensioni del batch. Pertanto, durante l'addestramento sulla TPU, imposta drop remainder=True per dataset.batch. Gli ultimi esempi di file potrebbero essere eliminati da un file per garantire che ogni batch abbia una forma statica batch_size. Ad esempio:

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)

Operazioni TensorFlow non disponibili

Messaggio di errore

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT
devices compatible with node

Framework e configurazioni interessati

Questo messaggio può essere visualizzato durante l'addestramento con TensorFlow.

Dettagli

Il modello utilizza un'operazione TensorFlow che non è attualmente disponibile sulla TPU.

Per un elenco delle operazioni disponibili sulla TPU, insieme ai piani per l'assistenza futura e i suggerimenti per le soluzioni alternative, consulta la guida alle operazioni di TensorFlow disponibili.

Messaggio di errore per esaurimento memoria

Messaggio di errore

ResourceExhaustedError: Ran out of memory in memory space hbm; used:
YYY; limit: 7.48G.

Framework e configurazioni interessati

Questo messaggio può essere visualizzato durante l'addestramento con TensorFlow, PyTorch o JAX.

Dettagli

Ogni Cloud TPU è composto da otto core TPU, le TPU v2 hanno 8 GB e le TPU v3 hanno 16 GB di RAM (o memoria HBM ad alta larghezza di banda). Questa memoria viene utilizzata per archiviare i tensori di ponderazione (variabile), nonché i tensori di risultato intermedi necessari per il calcolo dei gradienti. Se il modello è troppo grande per essere inserito nella RAM TPU, l'inizializzazione non riesce e il messaggio di errore riportato sopra viene stampato. Per ulteriori informazioni, consulta la sezione sulla riduzione dell'utilizzo della memoria.

Suggerimenti per ridurre l'utilizzo della memoria:

Problemi con l'interruzione dell'esecuzione

Se TensorFlow riscontra un errore durante l'esecuzione di TPU, lo script a volte sembra smettere di rispondere anziché uscire dalla shell. In questo caso, premi CTRL+\ sulla tastiera per attivare SIGQUIT, causando l'uscita immediata da parte di Python.

Allo stesso modo, la pressione di CTRL+C durante l'esecuzione di TPU non comporta l'arresto immediato di TensorFlow, ma attende invece fino alla fine del loop di iterazione corrente per l'uscita pulita.

Se si verificano nuovi errori durante la riconnessione alla TPU dopo l'uscita in questo modo, ripristina manualmente il server TPU con i comandi: 

gcloud compute tpus stop tpu-name --zone=zone
gcloud compute tpus start tpu-name --zone=zone

dove tpu-name viene ricavato dalla prima colonna visualizzata dal comando gcloud compute tpus list e zone è la zona mostrata nella seconda colonna.

Spaziatura interna eccessiva del tensore

Possibile causa di problema di memoria

I tensori in memoria TPU sono imbottiti, ovvero la TPU arrotonda le dimensioni dei tensori archiviati in memoria per eseguire i calcoli in modo più efficiente. Questa spaziatura interna avviene in modo trasparente a livello di hardware e non influisce sui risultati. In alcuni casi, tuttavia, la spaziatura interna può comportare un aumento significativo dell'utilizzo e dei tempi di esecuzione della memoria.

Come ridurre l'utilizzo di memoria

Il software TPU tenta di disporre i tensori in memoria per massimizzare l'efficienza di calcolo e ridurre al minimo la spaziatura interna. Questo processo di layout della memoria è tuttavia complesso, per ottenere risultati ottimali il modello deve rispettare la seguente regola generale. Per ridurre al minimo l'overhead di memoria e massimizzare l'efficienza di calcolo, è necessario che si verifichi una delle seguenti condizioni:

  • Le dimensioni totali del batch devono essere un multiplo di 64 (8 per core TPU) e le dimensioni della funzionalità devono essere un multiplo di 128.

    o

  • La dimensione totale del batch deve essere un multiplo di 1024 (128 per core TPU) e le dimensioni della funzionalità devono essere un multiplo di 8.

L'utilizzo di una dimensione batch di 1024 e dimensioni di funzionalità che sono un multiplo di 128 genera la migliore efficienza, anche se potrebbe non essere possibile per tutti i modelli. Per chiarezza, "dimensione della caratteristica" si riferisce alla dimensione nascosta di un livello connesso completamente o al numero di canali di output in una convoluzione. Non tutti i livelli possono essere conformi a questa regola, in particolare il primo e l'ultimo livello di rete. Non è un problema e si prevede che la maggior parte dei modelli richieda una certa spaziatura interna.

Riduzione dell'utilizzo della memoria

Se si verifica un errore di memoria esaurita durante l'esecuzione del modello sulla TPU, devi seguire alcuni passaggi per ridurre la memoria utilizzata dal modello.

I modi più efficaci per ridurre l'utilizzo di memoria sono:

  • Riduci la spaziatura interna con tensore eccessivo
  • Ridurre le dimensioni del batch

Dimensioni o modello del batch troppo grandi

Possibile causa di problema di memoria

Quando addestra una rete neurale su una CPU, GPU o TPU, l'utilizzo della memoria proviene da due punti:

  1. L'utilizzo della memoria è proporzionale al numero di ponderazioni del modello.
  2. Memorizzazione delle attivazioni intermedie dal passaggio in avanti necessario per calcolare il pass precedente. L'utilizzo della memoria è direttamente proporzionale alle dimensioni del batch, alle dimensioni dei livelli e al numero di livelli.

Pertanto, la memoria richiesta da un modello dipende in gran parte dalle dimensioni del batch.

La memoria richiesta da un modello dipende dal numero di livelli nella rete.

Il runtime TPU tenta di ottimizzare gli operatori per adattare il modello alla memoria (denominata rematerializzazione, simile al checkpoint di gradiente), ma non è sempre in grado di farlo.

Come ridurre l'utilizzo di memoria

Riduci lentamente le dimensioni del batch finché non entrano in memoria, assicurandoti che la dimensione totale del batch sia un multiplo di 64 (la dimensione del batch per core deve essere un multiplo di 8). Ricorda che i batch di dimensioni maggiori sono più efficienti sulla TPU. Una dimensione totale del batch pari a 1024 (128 per core) è in genere un buon punto di partenza.

Se il modello non può essere eseguito su TPU anche con una dimensione del batch di piccole dimensioni (ad esempio, 64), prova a ridurre il numero di livelli o le dimensioni del livello.

Migliorare la velocità di addestramento

Se il tuo modello è in grado di funzionare correttamente sulla TPU, ma la velocità di addestramento è inferiore al previsto, in questa sezione vengono descritti diversi modi per migliorare la velocità. Consulta la Guida al rendimento per altri suggerimenti su come migliorare il rendimento dell'addestramento.

Troppi pochi passaggi per esecuzione per ciclo di addestramento

Descrizione del problema di rendimento

Il trasferimento dell'argomento steps_per_execution a Model.compile determina il numero di passaggi di addestramento eseguiti tra i callback host. Ogni callback host richiede una comunicazione significativa tra la CPU host del server TPU e il dispositivo TPU, quindi se steps_per_execution è troppo piccolo, può rallentare l'addestramento.

Come sapere se il tuo modello è interessato

Se un profilo TPU rivela callback di CPU host frequenti tra i passaggi del dispositivo TPU, l'addestramento può trarre vantaggio da un valore steps_per_execution più elevato.

Come attenuare

Imposta steps_per_execution su un valore maggiore. Tieni presente che steps_per_execution può essere impostato su un valore elevato, ma tieni presente che la registrazione dei messaggi e il salvataggio di un checkpoint possono verificarsi solo dopo che è stato eseguito il numero di passaggi specificato.

Collo di bottiglia nell'elaborazione dell'input

Descrizione del problema di rendimento

Durante l'addestramento della TPU su un particolare blocco di dati, la funzione di elaborazione di input prepara il blocco di dati successivo sulla CPU. Se la funzione di input richiede più tempo rispetto alla funzione di modello, la TPU viene lasciata inattiva mentre la funzione di input recupera i dati.

Come sapere se il tuo modello è interessato

Segui le istruzioni nell'articolo Cloud TPU Tools: Input Pipeline Analyzer per visualizzare l'analisi della pipeline di input in TensorBoard:

immagine

La pagina di analisi della pipeline di input mostra un riepilogo chiaro che mostra se il tuo modello è stato bloccato da un'elaborazione di input. La stessa pagina mostra inoltre i tempi di esecuzione per operazione, che ti consentono di individuare le operazioni problematiche.

Come attenuare

Esistono diverse possibili mitigazioni durante il caricamento dei dati con l'API Dataset:

  1. Archivia i dati come una raccolta di strutture tf.train.Example in TFRecord file e caricali con TFRecordDataset. Consulta gli esempi di tutorial per set di dati o tutorial di ResNet.
  2. Utilizza dataset.cache() e/o dataset.prefetch() per eseguire il buffering dei dati di input. In questo modo si impedisce una sporadica lentezza nell'accesso ai file e la conseguente creazione di un collo di bottiglia.
  3. Specifica il parametro num_parallel_calls della funzione dataset.map() per abilitare le operazioni map() con più thread. Una semplice euristica per il valore di num_parallel_calls consiste nell'utilizzare il numero di core CPU disponibili.
  4. Esegui pre-elaborazione dei dati costosi offline come costo una tantum, invece di incorrere in costi in ogni periodo di addestramento.

L'intera elaborazione di input viene eseguita sulle CPU situate sul server TPU, non sulla macchina locale, quindi la velocità della macchina locale non è un fattore determinante.

Tempi di passo lenti e basso utilizzo di MXU

Descrizione del problema di rendimento

Cloud TPU è in grado di eseguire moltiplicazioni delle matrici e convoluzioni a velocità incredibilmente elevate. La maggior parte delle altre operazioni TensorFlow ha implementazioni efficienti sulla TPU, ma non sono la forza principale di TPU rispetto ad altri hardware. Pertanto, un modello dovrebbe essere dominato da moltiplicazioni delle matrici o convoluzioni per sfruttare appieno la TPU.

Come sapere se il tuo modello è interessato

I sintomi che troverai in questo caso sono tempi di passaggio lenti associati a un utilizzo ridotto dei record MXU mostrato quando imposti il profilo delle prestazioni.

Come attenuare

Prova a ridurre il numero di operazioni che non sono moltiplicazioni di matrici. Dopo aver ridotto il numero di moltiplicazioni delle matrici, ri-benchmark per verificare se le prestazioni sono accettabili sulle TPU.

Spaziatura interna eccessiva del tensore

Descrizione del problema di rendimento

I cuscinetti TPU si muovono in memoria in modo che la TPU possa utilizzare le sue unità di calcolo in modo efficiente. La spaziatura interna può aumentare l'utilizzo di memoria e larghezza di banda della memoria. Consulta la sezione Spaziatura interna tensore per comprendere e risolvere i problemi relativi alla spaziatura interna dei tensori.

Velocità effettiva lenta e utilizzo ridotto della memoria

Descrizione del problema di rendimento

Come regola generale, l'utilizzo di dimensioni batch più grandi determina una maggiore velocità di addestramento sulla TPU, in termini di esempi/secondo.

Come sapere se il tuo modello è interessato

La dimensione del batch di qualsiasi modello deve essere sempre almeno 64 (8 per core TPU), poiché la TPU riduce sempre i tensori a questa dimensione. Le dimensioni batch ideali per l'addestramento su TPU sono 1024 (128 per core TPU), poiché eliminano le inefficienze relative al trasferimento e alla spaziatura interna della memoria.

Come attenuare

La best practice prevede l'utilizzo delle dimensioni batch più grandi che si adattano alla memoria ed è un multiplo di 64. Il modo più semplice per raggiungere questo obiettivo è iniziare con 1024 e, se questo causa un errore di esaurimento della memoria, prova a ridurre le dimensioni del batch finché il modello non viene eseguito correttamente. La modifica delle dimensioni del batch di un modello può richiedere la regolazione di altri iperparametri per ottenere la stessa precisione del modello, ad esempio il tasso di apprendimento, ma questo deve essere valutato caso per caso.

Dimensioni del livello troppo piccole

Descrizione del problema di rendimento

Anche quando un modello è dominato da moltiplicazioni o convoluzioni delle matrici, la TPU potrebbe non funzionare alla massima efficienza se i tensori di input sono piccoli. Rispetto ad altri hardware, il TPU è più efficiente se le dimensioni del batch e delle dimensioni sono grandi (ad esempio, dimensione >= 512).

Come sapere se il tuo modello è interessato

Come regola generale, le dimensioni dei livelli inferiori a 128 raggiungono una scarsa efficienza sulla TPU, dato che 128 è la dimensione nativa dell'unità di moltiplicazione della matrice TPU. Per ottenere livelli ad alta efficienza per i livelli completamente connessi, si consiglia una dimensione nascosta minima 512. Ricorda che i livelli convoluzionali in genere non devono essere grandi quanto quelli completamente collegati per raggiungere lo stesso livello di efficienza.

Come attenuare

Se la motivazione principale per le taglie di strato piccole nel tuo modello è la velocità di addestramento, esegui nuovamente il benchmark dei tuoi modelli con strati più grandi sulla TPU. Ad esempio, aumentare le dimensioni dell'output di un livello da 256 a 512 può solo aumentare il tempo di addestramento del 20% anche se il modello esegue il doppio dei calcoli.

Profilazione di modelli a livello di opera d'arte

Spesso è utile misurare il tempo di esecuzione a livello di op e l'utilizzo della memoria per identificare i colli di bottiglia delle prestazioni. Per istruzioni su come eseguire questa operazione,
consulta la guida Strumenti di Cloud TPU: Visualizzatore tracce.

Eseguire il debug di cali dell'accuratezza del modello

Uno degli obiettivi dell'ecosistema Cloud TPU è che qualsiasi modello attualmente in fase di addestramento su una CPU o GPU ottiene una precisione molto simile quando viene addestrato sulla TPU, con forse piccoli aggiustamenti agli iperparametri come la dimensione del batch e la velocità di apprendimento. A volte, tuttavia, gli utenti possono notare un degrado di accuratezza durante l'addestramento dei modelli sulla TPU. Il debug di questi problemi può essere estremamente frustrante a causa della natura casuale dell'addestramento della rete neurale. Questa sezione fornisce indicazioni su come individuare la causa principale di eventuali cali di accuratezza del modello quando trasferisci un modello alla TPU.

Informazioni sullo sharding dei dati (parallelism dei dati)

Uno degli obiettivi principali di TensorFlow è che ogni operazione deve produrre risultati quasi identici, indipendentemente dal fatto che venga eseguita su CPU, GPU o TPU. Esistono alcune eccezioni a questa regola, ad esempio le operazioni casuali. In generale, se rilevi una differenza significativa tra l'output delle operazioni non casuali sulla TPU e sulla CPU, segnalalo come bug.

Tuttavia, per la pipeline di addestramento nel suo insieme, esiste una significativa differenza tra l'addestramento su CPU/GPU e TPU. Quando si esegue l'addestramento su una TPU, TensorFlow esegue lo sharding dei dati, Ogni Cloud TPU contiene 8 core TPU che operano come unità di elaborazione indipendenti. Per ogni passaggio dell'addestramento, ogni core TPU riceve un batch di dati, calcola i gradienti di peso, scambia i gradienti con gli altri core TPU, quindi calcola l'aggiornamento del peso. Per impostazione predefinita, la perdita viene calcolata in media tra i core, ma può essere sommata modificando il parametro CrossShardOptimizer.

Se la perdita totale del modello può essere calcolata come la media (o somma) delle perdite indipendenti per campione, questa procedura è matematicamente equivalente all'addestramento su un singolo batch di grandi dimensioni.

L'operatività più comune non indipendente per campione è la normalizzazione batch, che viene eseguita separatamente su ciascun batch per core. Ad esempio, se la dimensione totale del batch è 128, la dimensione del batch per core è 16 e ciascuno degli 8 core esegue la normalizzazione sui propri 16 campioni. In alcuni casi, la normalizzazione dei batch su piccoli batch (ad esempio, meno di 32) ha riscontrato un peggioramento della precisione. Nello scenario ideale, le dimensioni totali del batch devono essere grandi (ad esempio, da 256 a 1024). Se la dimensione del batch è troppo grande per essere inserita in memoria, l'effetto dello sharding deve essere valutato caso per caso.

Addestramento determinazione

Un motivo per cui è difficile eseguire il debug delle differenze di accuratezza del modello è che in diversi framework (TensorFlow, PyTorch, JAX), il software di addestramento utilizza diverse inizializzazioni di peso e data shuffling ogni volta che un modello viene addestrato. È bene modificare la procedura di addestramento in modo che sia deterministica, in modo che più esecuzioni restituiscano modelli quasi identici. Questa sezione illustra come eseguire il tutorial MNIST in modo deterministico:

  1. Generare un file di checkpoint iniziale eseguendo un singolo passaggio sulla CPU. Il passaggio viene utilizzato per ottenere l'inizializzazione del peso deterministico. Inoltre, assicurati di utilizzare un seed casuale fisso per qualsiasi funzione casuale del modello. 
# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1
  1. Modifica le funzioni di data shuffling nella funzione di input per utilizzare un seed casuale. Questo è già stato fatto nel tutorial MNIST. Questo vale per le operazioni di elaborazione dei dati di input, perché vengono sempre eseguite sulla CPU. Le operazioni casuali nella funzione del modello potrebbero non essere deterministiche tra la TPU e la CPU. Ogni volta che chiami un'operazione casuale, passa un seed fisso per garantire gli stessi risultati tra le esecuzioni. Ad esempio: 
# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)
  1. Esegui lo stesso modello due volte sulla CPU per verificare che l'addestramento sia deterministico. Tieni presente che l'addestramento deve essere eseguito per un numero ragionevole di passaggi (ad esempio 1000), ma non deve essere eseguito per raggiungere la convergenza.

    Dato che l'addestramento della CPU viene confrontato con un addestramento di un singolo core TPU, utilizza una dimensione del batch che può adattarsi a un singolo core TPU (in genere, la dimensione intera del batch diviso per 8). TensorFlow non garantisce il determinismo bit per bit tra le esecuzioni, ma la perdita dovrebbe essere molto simile:

Copia le ponderazioni iniziali

gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

Gara 1

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

Output 1

accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

Gara 2

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

Uscita 2

accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

Addestramento TPU singolo core

Dopo aver eseguito il tutorial MNIST in modo deterministico, il passaggio successivo consiste nella replica dei risultati addestrati dalla CPU sulla TPU, utilizzando un singolo core TPU per indicare se il problema è legato allo sharding dei dati o al motore di esecuzione della TPU stesso.

Ecco come eseguire l'addestramento e la valutazione single-core sul tutorial MNIST:

Utilizza la stessa inizializzazione sulla ponderazione della CPU.

gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

Esegui l'addestramento per 1000 passi

python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --num_shards=1 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

Output

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

La perdita non corrisponderà esattamente al modello addestrato dalla CPU, ma dovrebbe essere vicina. Se non è simile al tuo modello, potrebbe significare che hai trovato un bug nel motore di esecuzione della TPU. Prima di inviare una segnalazione di bug, verifica quanto segue:

  1. Stai passando num_shards=1 a TPUConfig.

  2. Non esistono operazioni casuali nella funzione del modello e le operazioni casuali nella funzione di input vengono sottoposte a seeding corretto.

  3. Stai utilizzando lo stesso file di checkpoint iniziale per l'addestramento della CPU e della TPU.

Debug dell'addestramento della TPU multi-core

Se il tuo modello ottiene la stessa perdita sulla CPU e sul TPU single-core, il problema è probabilmente uno dei seguenti:

(a) Il degrado è dovuto alla varianza naturale naturale durante l'addestramento di modelli neurali con inizializzazioni diverse.

(b) Il degrado è dovuto a un problema relativo allo sharding dei dati sulla TPU.

Per determinare se (a) è il problema, riaddestra il modello completo su CPU/GPU e TPU multi-core utilizzando la stessa inizializzazione sul peso, come sopra.

Se hai la certezza che il calo di accuratezza sia statisticamente significativo, i problemi più probabili relativi allo sharding dei dati sono:

  1. Se il tuo modello utilizza la normalizzazione batch, una dimensione totale del batch inferiore a 256 (ad esempio, inferiore a 32 per core) potrebbe ridurre la precisione.
  2. Le funzioni di perdita in modalità batch sono interessate dallo sharding. Queste funzioni di perdita sono in genere piuttosto specializzate. Ad esempio, Karras et al.2017 utilizza un discriminatore batch durante l'addestramento di una rete generativa di contraddizione.

Risoluzione dei problemi delle VM TPU

I seguenti problemi e soluzioni sono applicabili solo alle configurazioni di VM TPU.

Risoluzione dei problemi di configurazione di gcloud

Problema
gcloud components update visualizza il seguente messaggio di errore:
ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager is
disabled for this installation.
Soluzione
Per utilizzare gcloud con una VM TPU, devi utilizzare un'installazione gcloud non gestita tramite un gestore di pacchetti. Per installare gcloud dal codice sorgente:
  sudo apt-get remove google-cloud-sdk
  curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  ./google-cloud-sdk/install.sh
  source ~/.bashrc
Problema

Il comando gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} mostra il seguente messaggio di errore:

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.
Soluzione

Potrebbe essersi verificato un problema di propagazione della chiave SSH. Prova a spostare le chiavi generate automaticamente in una posizione di backup per forzare gcloud a ricrearle:

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

Log di debug

I framework Cloud TPU supportati, JAX, PyTorch e TensorFlow accedono alle TPU tramite una libreria condivisa denominata libtpu presente in ogni VM TPU. Questa libreria include il compilatore XLA utilizzato per compilare i programmi TPU, il runtime TPU utilizzato per eseguire i programmi compilati e il driver TPU utilizzato dal runtime per l'accesso di basso livello alla TPU.

La libreria libtpu registra informazioni che possono essere utili per il debug. Per impostazione predefinita, questi log vengono scritti in /tmp/tpu_logs su ogni VM Cloud TPU. Puoi impostare le seguenti variabili di ambiente prima di iniziare l'addestramento per modificare il comportamento di logging:

TPU_LOG_DIR: la directory in cui vengono scritti i log
Il percorso predefinito della directory è /tmp/tpu_logs. La directory viene creata se non esiste già, ma non vengono create directory padre. Se si verifica un errore durante la ricerca o la creazione della directory specificata, un messaggio viene stampato su stderr, ma non interrompe il programma e il logging viene disabilitato. Imposta il nome della directory su "disattivato" per disabilitare del tutto il logging sul disco.
TPU_MIN_LOG_LEVEL: la gravità minima che verrà registrata sul disco
Le opzioni sono 0 (INFO), 1 (AVVISO), 2 (ERRORE) e 3 (FATALE). Il valore predefinito è 0.
TPU_STDERR_LOG_LEVEL: la gravità minima che verrà registrata su stderr, oltre al disco, se applicabile
Le opzioni disponibili sono le stesse di TPU_MIN_LOG_LEVEL. Il valore predefinito è 3.
TPU_MAX_LOG_SIZE_MB: la dimensione massima in megabyte di ogni file di log
Un nuovo file di log verrà avviato automaticamente quando quello precedente raggiungerà approssimativamente le dimensioni. Il valore predefinito è 1024.