Grazie per esserti sintonizzato su Google I/O. Visualizza tutte le sessioni su richiesta Guarda su richiesta

XLA: ottimizzazione del compilatore per l'apprendimento automatico

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 TensorFlow. Ogni operazione TensorFlow ha un'implementazione del kernel GPU precompilata a cui l'executor 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 univoci per il modello, possono sfruttare le informazioni specifiche del modello per l'ottimizzazione. Ad esempio, diamo un'occhiata a un'ottimizzazione che XLA esegue 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 avvia tre kernel: uno per la moltiplicazione, uno per l'addizione e uno per la riduzione. Tuttavia, XLA può ottimizzare il grafico in modo da calcolare il risultato in un singolo avvio del kernel. Lo fa "fondendo" l'addizione, la moltiplicazione e la riduzione in un singolo kernel GPU. Inoltre, questa operazione fusa non scrive nella memoria i valori intermedi prodotti da y*z e x+y*z ; invece "trasmette" i risultati di questi calcoli intermedi direttamente ai propri utenti mantenendoli interamente nei registri della GPU. Fusion è l'ottimizzazione singola 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 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 di compilazione obbligatoria : l'intera funzione viene compilata con XLA oppure viene generata un'eccezione errors.InvalidArgumentError . XLA non può attualmente compilare funzioni in cui le dimensioni non sono inferibili : ovvero, 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 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]))

Guarda il tutorial colab per un esempio di utilizzo più dettagliato e un video tutorial sull'utilizzo 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 strategia distribuita

XLA:GPU può essere utilizzato con la strategia distribuita TF ( MirroredStrategy o MultiWorkerMirroredStrategy ) annotando la funzione passo 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 alcuna modifica è abilitare l'auto-clustering , che trova automaticamente i cluster (sottografi connessi) all'interno delle funzioni TensorFlow che possono essere compilati ed eseguiti utilizzando XLA. Il clustering automatico su GPU può essere abilitato impostando la variabile di 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 della GPU, ma può essere abilitato anche 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 l' esercitazione sul clustering automatico colab .

Compilazione AOT (Ahead-of-time) per CPU con tfcompile

Puoi anche utilizzare uno strumento tfcompile autonomo, che converte il grafico TensorFlow in codice eseguibile (solo per CPU x86-64).

Ispezionare 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 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.txt Programmi XLA generati, uno per ogni cluster compilato. Allegarli quando si inviano segnalazioni di bug XLA è estremamente utile!

  • module_XXXX.ir-*.ll File generati nella rappresentazione intermedia LLVM , con intrinseche NVPTX .

  • module_XXXX.ptx File PTX generati.

Puoi anche eseguire il dump del grafico visualizzando l'incorporamento dei 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

Una segnalazione di bug è molto più facile da riprodurre se include i dump per i programmi XLA generati e l'incorporamento del clustering automatico utilizzato. Per generarli per un programma TensorFlow in esecuzione con 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 archiviano i bug, allegare il contenuto della directory /tmp/generated (indicata sopra).

Se possibile, prova a isolare un bug in un singolo programma XLA utilizzando run_hlo_module ed eseguendolo in modo iterativo sui programmi generati.

Ulteriori letture

Frontend XLA

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

  • JAX : trasformazioni componibili di programmi Python+NumPy
  • Julia : il linguaggio Julia per il calcolo scientifico
  • PyTorch : quadro PyTorch
  • Nx : libreria di calcolo numerico per il linguaggio di programmazione Elixir

Colloqui

Usando XLA da TF usando jit_compile=True

Panoramica XLA