XLA: optimiser le compilateur pour le machine learning

OpenXLA est un compilateur d'algèbre linéaire spécifique à un domaine qui permet d'accélérer les modèles TensorFlow sans qu'il faille nécessairement modifier le code source.

Introduction

Lorsqu'un programme TensorFlow est exécuté, l'exécuteur TensorFlow traite toutes les opérations séparément. Chaque opération TensorFlow comprend une implémentation de noyau GPU précompilée vers laquelle renvoie l'exécuteur.

XLA propose un autre mode d'exécution des modèles: il compile le graphe TensorFlow dans une séquence de noyaux de calcul générés spécialement pour un modèle en particulier. Étant donné que ces noyaux sont spécifiques au modèle, ils peuvent exploiter des informations qui lui sont propres en vue de l'optimisation. Prenons l'exemple d'une optimisation que XLA effectue dans le cadre d'un calcul TensorFlow simple:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

Exécuté sans XLA, le graphe lance trois noyaux: un pour la multiplication, un pour l'addition et un pour la réduction. Toutefois, XLA peut optimiser le graphe de sorte qu'il calcule le résultat dans un seul lancement de noyau. Pour ce faire, il fusionne les opérations d'addition, de multiplication et de réduction dans un seul noyau GPU. De plus, cette opération fusionnée n'écrit pas en mémoire les valeurs intermédiaires générées par y*z et x+y*z. En revanche, elle "diffuse" les résultats de ces calculs intermédiaires directement aux utilisateurs, tout en les conservant intégralement dans des registres GPU. La fusion est l'optimisation la plus importante de XLA. La bande passante mémoire est généralement la ressource la plus limitée des accélérateurs matériels. Par conséquent, la suppression des opérations de mémoire constitue l'un des meilleurs moyens d'améliorer les performances.

Activer XLA pour les modèles TensorFlow

Compilation explicite avec tf.function(jit_compile=True)

L'API de compilation explicite permet de choisir avec précision les fonctions à compiler. Par exemple, la fonction TensorFlow suivante, qui effectue l'entraînement de l'ensemble de données MNIST, est compilée avec 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 présente une sémantique must-compile: soit toute la fonction est compilée avec XLA, soit une exception errors.InvalidArgumentError est générée. Pour l'heure, XLA ne peut pas compiler les fonctions dont les dimensions ne peuvent pas être déduites: en d'autres termes, il est impossible de déduire les dimensions de tous les Tensors sans exécuter l'intégralité du calcul. Par exemple, la fonction suivante ne sera pas compilée:

@tf.function
def not_compilable(x):
  return tf.unique(x)

Cependant, les formes peuvent varier d'une exécution à l'autre:

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

Consultez ce tutoriel au format Colab pour un exemple plus détaillé, et ce tutoriel vidéo pour comprendre comment utiliser jit_compile=True.

Utilisation avec Keras

Pour les modèles Keras, jit_compile=True peut être défini comme argument de model.compile:

model.compile(optimizer="adam", jit_compile=True)

Utilisation avec une stratégie distribuée

XLA:GPU peut être utilisé avec la stratégie distribuée TF (MirroredStrategy ou MultiWorkerMirroredStrategy) en annotant la fonction de pas avec 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)

Clustering automatique

Pour commencer à utiliser XLA dans les modèles TensorFlow sans apporter la moindre modification, une méthode simple consiste à activer le clustering automatique. Cette opération recherche automatiquement dans les fonctions TensorFlow les clusters (sous-graphes associés) qui peuvent être compilés et exécutés à l'aide de XLA. Le clustering automatique sur GPU peut être activé en définissant la variable d'environnement TF_XLA_FLAGS:

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

Actuellement, le clustering automatique est optimisé pour les charges de travail de GPU, mais il peut également être activé sur le processeur en utilisant, en outre, l'indicateur --tf_xla_cpu_global_jit:

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

Pour obtenir un exemple d'utilisation détaillé, consultez le tutoriel Colab relatif au clustering automatique.

Compilation anticipée (AOT) pour le processeur avec tfcompile

Vous pouvez également utiliser un outil tfcompile autonome qui convertit le graphe TensorFlow en code exécutable (pour un processeur x86-64 uniquement).

Inspecter les programmes compilés

XLA fournit des fonctionnalités d'introspection qui vous permettent d'inspecter les programmes générés. Pour vider les programmes générés, utilisez la variable d'environnement XLA_FLAGS:

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

Une fois le vidage effectué, vous trouverez les fichiers suivants dans /tmp/generated:

  • module_XXXX.*_optimizations.txt Programmes XLA générés, un par cluster compilé. Il est vivement conseillé de joindre ces programmes lors de l'envoi de rapports de bugs XLA.

  • module_XXXX.ir-*.ll Fichiers générés dans la représentation intermédiaire LLVM, avec des fonctionnalités intrinsèques NVPTX.

  • module_XXXX.ptx Fichiers PTX générés.

Vous pouvez également vider le graphe qui visualise l'imbrication des clusters XLA dans le graphe TensorFlow avec la ligne de code suivante:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

Rapports de bugs reproductibles

Un rapport de bug est beaucoup plus facile à reproduire s'il inclut des vidages pour les programmes XLA générés et la représentation vectorielle utilisée pour le clustering automatique. Si vous souhaitez les générer pour un programme TensorFlow exécuté avec le clustering automatique, lancez:

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

Lorsque vous signalez des bugs, joignez le contenu du répertoire /tmp/generated (référencé ci-dessus).

Si possible, essayez d'isoler un bug dans un seul programme XLA en utilisant l'outil run_hlo_module et en l'exécutant de manière itérative sur les programmes générés.

Documentation complémentaire

Interfaces XLA

Outre TensorFlow, les programmes XLA peuvent être générés par:

  • JAX: transformations composables de programmes Python+NumPy
  • Julia: langage Julia utilisé pour le calcul scientifique
  • PyTorch: framework PyTorch
  • Nx: bibliothèque de calcul numérique pour le langage de programmation Elixir

Conférences

Utiliser XLA à partir de TF avec jit_compile=True

Présentation de XLA