XLA: compilatore di ottimizzazione per l'apprendimento automatico

XLA (Accelerated Linear Algebra) è un compilatore specifico del dominio per l'algebra lineare in grado di accelerare i modelli TensorFlow potenzialmente senza modifiche al codice sorgente.

I risultati sono miglioramenti nella velocità e della memoria: ad esempio in BERT MLPerf presentazione utilizzando 8 GPU Volta V100 utilizzando XLA ha raggiunto un miglioramento delle prestazioni ~ 7x e il miglioramento dimensione ~ 5x batch:

introduzione

Quando viene eseguito un programma TensorFlow, tutte le operazioni vengono eseguite individualmente dall'esecutore TensorFlow. Ogni operazione TensorFlow ha un'implementazione del kernel GPU precompilata a cui invia l'executor.

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 che XLA fa 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 avvio del kernel. Lo fa "fondendo" l'aggiunta, la moltiplicazione e la riduzione in un singolo kernel GPU. Inoltre, questa operazione fusa non scrive i valori intermedi prodotti da y*z e x+y*z alla memoria; invece "trasmette" i risultati di questi calcoli intermedi direttamente ai propri utenti mantenendoli interamente nei registri GPU. Fusion è l'ottimizzazione più importante di XLA. La larghezza di banda della memoria è in genere la risorsa più scarsa sugli acceleratori hardware, quindi la rimozione delle 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 granulare per la scelta delle funzioni da compilare. Ad esempio, la seguente funzione TensorFlow che esegue la formazione 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))

Il jit_compile API contiene semantica must-compilazione: o l'intera funzione viene compilato con XLA, o un errors.InvalidArgumentError viene generata un'eccezione. XLA attualmente non può compilare funzioni cui dimensioni non sono corollari: che è, se non è possibile dedurre le dimensioni di tutti 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 durante le corse 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]))

Vedere la CoLab esercitazione per un esempio più dettagliato l'utilizzo, e un video tutorial su jit_compile=True utilizzo.

Clustering automatico

Un modo semplice per iniziare a utilizzare XLA in modelli tensorflow senza alcuna modifica è consentire l'auto-clustering, che trova automaticamente cluster (sottografi collegati) all'interno delle funzioni tensorflow che può essere compilato ed eseguito utilizzando XLA. Auto-il clustering su GPU può essere abilitata impostando la TF_XLA_FLAGS variabile d'ambiente:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Auto-cluster è attualmente ottimizzato per carichi di lavoro GPU, ma può anche essere attivata CPU ulteriormente 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 vedere l' auto-clustering di esercitazione CoLab .

AOT (ahead-of-time) di compilazione per CPU con tfcompile

È anche possibile utilizzare un autonomo tfcompile strumento, che converte il grafico tensorflow in codice eseguibile (solo CPU x86-64).

Ispeziona i programmi compilati

XLA fornisce funzionalità 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 la viene eseguita al dumping, è possibile trovare i seguenti file in /tmp/generated :

  • module_XXXX.*_optimizations.txt generati programmi XLA , uno per ogni cluster compilato. Allegare quelli quando si inviano segnalazioni di bug XLA è estremamente utile!

  • module_XXXX.ir-*.ll file generati in LLVM rappresentazione intermedia, con NVPTX intrinseche.

  • module_XXXX.ptx generata PTX file.

Puoi anche scaricare il grafico visualizzando 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 clustering automatico, avviare:

$ 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"

Al momento del deposito insetti, fissare il contenuto della /tmp/generated directory (di cui sopra).

Se possibile, cercare di isolare un bug per un singolo programma XLA utilizzando il replay_computation e iterativo in esecuzione su programmi generati.

Ulteriori letture

Frontend XLA

Oltre a TensorFlow, i programmi XLA possono essere generati da:

  • JAX trasformazioni componibili di programmi Python + numpy:
  • Julia : La lingua Julia per il calcolo scientifico
  • PyTorch quadro PyTorch:
  • Nx : Numerical library di calcolo per il linguaggio di programmazione Elixir

colloqui

Utilizzando XLA da TF utilizzando jit_compile=True

Panoramica XLA