XLA: ottimizzazione del compilatore per l'apprendimento automatico

Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

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

I risultati sono miglioramenti nella velocità e nell'utilizzo della memoria: ad esempio, nell'invio BERT MLPerf utilizzando 8 GPU Volta V100 utilizzando XLA ha ottenuto un miglioramento delle prestazioni di ~7 volte e un miglioramento della dimensione del batch ~5 volte superiore:

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'esecutore esegue la spedizione.

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 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 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'addizione, la moltiplicazione e la riduzione in un unico kernel GPU. Inoltre, questa operazione fusa non scrive in memoria i valori intermedi prodotti da y*z e 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 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 dettagliato per la scelta delle funzioni da compilare. Ad esempio, la seguente funzione TensorFlow che esegue il training 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 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 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)

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]))

Vedere il tutorial colab per un esempio di utilizzo più dettagliato e un video tutorial su jit_compile=True usage.

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 step 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 collegati) all'interno delle funzioni TensorFlow che possono essere compilate ed eseguite 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

L'auto-clustering è attualmente ottimizzato per i carichi di lavoro della GPU, ma può essere abilitato anche sulla CPU utilizzando inoltre 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

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 elementi intrinseci NVPTX .

  • module_XXXX.ptx File PTX generati.

Puoi anche scaricare il 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 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"

Quando si archiviano i bug, allegare il contenuto della /tmp/generated (riferita sopra).

Se possibile, prova a isolare un bug in un singolo programma XLA utilizzando replay_computation 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 l'informatica scientifica
  • PyTorch : framework PyTorch
  • Nx : Libreria di calcolo numerico per il linguaggio di programmazione Elixir

Colloqui

Utilizzo di XLA da TF utilizzando jit_compile=True

Panoramica di XLA