OpenXLA è un compilatore specifico per dominio per l'algebra lineare che può accelerare i modelli TensorFlow senza apportare modifiche al codice sorgente.
Introduzione
Quando viene eseguito un programma TensorFlow, tutte le operazioni vengono eseguite individualmente dall'esecutore TensorFlow. Ogni operazione TensorFlow ha un'implementazione del kernel della GPU precompilata a cui l'esecutore invia.
XLA fornisce una modalità alternativa di esecuzione dei modelli: compila il grafo di TensorFlow in una sequenza di kernel di calcolo generati specificamente per il determinato modello. Poiché questi kernel sono specifici del modello, possono sfruttare informazioni specifiche del modello per l'ottimizzazione. Ad esempio, diamo un'occhiata a un'optimizzazione eseguita da XLA nel contesto di un semplice calcolo TensorFlow:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
Se viene eseguito senza XLA, il grafico avvia tre kernel: uno per la moltiplicazione, uno per l'addizione e uno per la riduzione. Tuttavia, XLA può ottimizzare il gráfo in modo da calcolare il risultato in un singolo avvio del kernel. Per farlo, "fonde" l'addizione, la moltiplicazione e la riduzione in un singolo kernel della GPU.
Inoltre, questa operazione combinata non scrive i valori intermedi
prodotti da y*z e x+y*z nella memoria, ma "trasmette" i risultati di
questi calcoli intermedi direttamente ai relativi utenti, mantenendoli
interamente nei registri della GPU. La fusione è l'ottimizzazione più importante di XLA.
La larghezza di banda della memoria è in genere la risorsa più scarsa negli acceleratori hardware, quindi rimuovere le operazioni di memoria è uno dei modi migliori per migliorare le prestazioni.
Attivare XLA per i modelli TensorFlow
Compilazione esplicita con tf.function(jit_compile=True)
L'API di compilazione esplicita offre un controllo granulare per scegliere quali funzioni compilare. Ad esempio, la seguente funzione TensorFlow che esegue l'addestramento di MNIST viene 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 must-compile: l'intera funzione viene compilata con XLA o viene sollevata un'eccezione errors.InvalidArgumentError. Al momento XLA non può compilare funzioni in cui le dimensioni non sono
deducibili: ovvero se non è possibile dedurre 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)
Tuttavia, le forme possono variare da una corsa all'altra:
@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]))
Consulta il tutorial di Colab per un esempio di utilizzo più dettagliato e un video tutorial sull'utilizzo di jit_compile=True.
Utilizzo con Keras
Per i modelli Keras, jit_compile=True può essere impostato come argomento per
model.compile:
model.compile(optimizer="adam", jit_compile=True)
Utilizzo con la strategia distribuita
XLA:GPU può essere utilizzata con la strategia distribuita di TF
(MirroredStrategy
o
MultiWorkerMirroredStrategy)
annotando la funzione di passaggio con jit_compile=True:
@tf.function(jit_compile=True)
def step_fn():
t = tf.ones(shape=[100], dtype=tf.float32)
ctx = tf.distribute.get_replica_context()
return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)
@tf.function
def run_fn():
return strategy.run(step_fn)
Raggruppamento automatico
Un modo semplice per iniziare a utilizzare XLA nei modelli TensorFlow senza apportare modifiche è attivare l'aggregazione automatica, che trova automaticamente cluster (sottografi collegati) all'interno delle funzioni TensorFlow che possono essere compilati ed eseguiti utilizzando XLA. Il clustering automatico su GPU può essere attivato impostando la variabile di ambiente TF_XLA_FLAGS:
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
L'aggregazione automatica è attualmente ottimizzata per i workload GPU, ma può essere attivata anche su CPU utilizzando 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 dettagliato di utilizzo, consulta il tutorial di Colab sull'autoclustering.
Compilazione AOT (Ahead-of-time) per CPU con tfcompile
Puoi anche utilizzare uno strumento autonomo tfcompile, che converte il grafo TensorFlow in codice eseguibile (solo per CPU x86-64).
Controllare i programmi compilati
XLA fornisce funzionalità di introspezione che ti consentono di ispezionare i programmi generati. Per eseguire il dump dei programmi generati, utilizza la variabile di 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, puoi trovare i seguenti file in
/tmp/generated:
module_XXXX.*_optimizations.txtProgrammi XLA generati, uno per ogni cluster compilato. Allegarli quando invii segnalazioni di bug XLA è estremamente utile.module_XXXX.ir-*.llFile generati in rappresentazione intermedia LLVM, con istruzioni intrinsic NVPTX.module_XXXX.ptxFile PTX generati.
Puoi anche eseguire il dump del grafico che visualizza l'embedding dei cluster XLA all'interno del grafico di TensorFlow con:
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
Segnalazioni di bug riproducibili
È molto più facile riprodurre un report di bug se include i dump per i programmi XLA generati e l'embedding di clustering 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 invii i bug, allega i contenuti della directory /tmp/generated
(a cui si fa riferimento sopra).
Se possibile, prova a isolare un bug in un singolo programma XLA utilizzando run_hlo_module e eseguendolo in modo iterativo sui programmi generati.
Per approfondire
- Documentazione di OpenXLA Documentazione di OpenXLA
- Problemi noti Elenco dei problemi noti relativi a XLA+TF
- XLA - TensorFlow, compilato: leggi il post del blog Google Developers
- Dai un'occhiata al codice sorgente XLA su GitHub.
Frontend XLA
Oltre a TensorFlow, i programmi XLA possono essere generati da:
- JAX: trasformazioni componibili di programmi Python+NumPy
- Julia: il linguaggio Julia per la calcolo scientifico
- PyTorch: framework PyTorch
- Nx: libreria di calcolo numerico per il linguaggio di programmazione Elixir