XLA (Accelerated Linear Algebra) è un compilatore specifico del dominio per l'algebra lineare che può accelerare i modelli TensorFlow potenzialmente senza modifiche al codice sorgente.
I risultati sono miglioramenti nella velocità e nell'utilizzo della memoria: ad esempio, l' invio di BERT MLPerf utilizzando 8 GPU Volta V100 utilizzando XLA ha ottenuto un miglioramento delle prestazioni di ~ 7 volte e un miglioramento delle dimensioni del batch di ~ 5 volte:
introduzione
Quando viene eseguito un programma TensorFlow, tutte le operazioni vengono eseguite individualmente dall'esecutore di TensorFlow. Ogni operazione di TensorFlow ha un'implementazione del kernel GPU precompilata a cui l'esecutore invia.
XLA fornisce una modalità alternativa di esecuzione dei modelli: compila il grafico TensorFlow in una sequenza di kernel di calcolo generati specificamente per il modello dato. Poiché questi kernel sono unici per il modello, possono sfruttare le informazioni specifiche del modello per l'ottimizzazione. Ad esempio, diamo un'occhiata a un'ottimizzazione eseguita da XLA nel contesto di un semplice calcolo TensorFlow:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
Eseguito senza XLA, il grafico lancia tre kernel: uno per la moltiplicazione, uno per l'addizione e uno per la riduzione. Tuttavia, XLA può ottimizzare il grafico in modo che calcoli il risultato in un singolo lancio del kernel. Lo fa "fondendo" l'aggiunta, la moltiplicazione e la riduzione in un unico kernel GPU. Inoltre, questa operazione di fusione non scrive in memoria i valori intermedi prodotti da y*z x+y*z ; invece "trasmette" i risultati di questi calcoli intermedi direttamente ai loro utenti mantenendoli interamente nei registri della GPU. Fusion è l'ottimizzazione più importante di XLA. La larghezza di banda della memoria è in genere la risorsa più scarsa sugli acceleratori hardware, quindi rimuovere le operazioni di memoria è uno dei modi migliori per migliorare le prestazioni.
Abilita XLA per i modelli TensorFlow
Compilazione esplicita con tf.function(jit_compile=True)
L'API di compilazione esplicita offre un controllo dettagliato per la scelta delle funzioni da compilare. Ad esempio, la seguente funzione TensorFlow che esegue l'addestramento MNIST è compilata con XLA:
@tf.function(jit_compile=True)
def train_mnist(images, labels):
images, labels = cast(images, labels)
with tf.GradientTape() as tape:
predicted_labels = layer(images)
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=predicted_labels, labels=labels
))
layer_variables = layer.trainable_variables
grads = tape.gradient(loss, layer_variables)
optimizer.apply_gradients(zip(grads, layer_variables))
L'API jit_compile ha una semantica da compilare : l'intera funzione viene compilata con XLA o viene generata un'eccezione errors.InvalidArgumentError . XLA non può attualmente compilare funzioni in cui le dimensioni non sono inferibili : cioè, se non è possibile inferire le dimensioni di tutti i tensori senza eseguire l'intero calcolo. Ad esempio, la seguente funzione non verrà compilata:
@tf.function
def not_compilable(x):
return tf.unique(x)
Le forme possono variare tra le piste però:
@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
return a + b
recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))
Vedi il tutorial colab per un esempio di utilizzo più dettagliato.
Auto-clustering
Un modo semplice per iniziare a utilizzare XLA nei modelli TensorFlow senza alcuna modifica è abilitare il clustering automatico, che trova automaticamente i cluster (sottografi connessi) all'interno delle funzioni TensorFlow che possono essere compilati ed eseguiti utilizzando XLA. Il clustering automatico sulla GPU può essere abilitato impostando la variabile d'ambiente TF_XLA_FLAGS :
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
Il clustering automatico è attualmente ottimizzato per i carichi di lavoro GPU, ma può anche essere abilitato sulla CPU utilizzando anche il flag --tf_xla_cpu_global_jit :
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
Per un esempio di utilizzo dettagliato, vedere il tutorial sul clustering automatico colab .
Compilazione AOT (Ahead-of-time) per CPU con tfcompile
È inoltre possibile utilizzare uno strumento tfcompile autonomo, che converte il grafico TensorFlow in codice eseguibile (solo per CPU x86-64).
Ispeziona i programmi compilati
XLA fornisce servizi di introspezione che consentono di ispezionare i programmi generati. Per eseguire il dump dei programmi generati, utilizzare la variabile d'ambiente XLA_FLAGS :
$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
Dopo aver eseguito il dumping, è possibile trovare i seguenti file in /tmp/generated :
module_XXXX.*_optimizations.txtGenerato programmi XLA , uno per ogni cluster compilato. Allegarli quando si inviano segnalazioni di bug XLA è estremamente utile!module_XXXX.ir-*.llFile generati nella rappresentazione intermedia LLVM , con intrinseci NVPTX .module_XXXX.ptxFile PTX generati.
Puoi anche scaricare il grafico che visualizza l'incorporamento di cluster XLA all'interno del grafico TensorFlow con:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
Segnalazioni di bug riproducibili
Un bug report è molto più facile da riprodurre se include dump per i programmi XLA generati e l'incorporamento di cluster automatico utilizzato. Per generarli per un programma TensorFlow in esecuzione con il clustering automatico, avvia:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
my/tensorflow/program"
Quando si segnalano bug, allegare il contenuto della /tmp/generated (a cui si fa riferimento sopra).
Se possibile, prova a isolare un bug in un singolo programma XLA utilizzando replay_computation ed eseguendolo in modo iterativo sui programmi generati.
Ulteriore lettura
- Problemi noti Elenco dei problemi noti con XLA
- Architettura XLA : panoramica dell'architettura XLA
- XLA - TensorFlow, compilato : leggi sul blog degli sviluppatori di Google
- Controlla la fonte XLA su Github!
Frontend XLA
Oltre a TensorFlow, i programmi XLA possono essere generati da: