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 avoir à modifier le code source.

Introduction

Lorsqu'un programme TensorFlow est exécuté, toutes les opérations sont exécutées individuellement par l'exécuteur TensorFlow. Chaque opération TensorFlow comporte une implémentation de noyau GPU précompilée vers laquelle l'exécuteur est envoyé.

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écifiquement pour le modèle donné. Étant donné que ces noyaux sont propres au modèle, ils peuvent exploiter des informations spécifiques au modèle à des fins d'optimisation. Examinons par exemple une optimisation effectuée par XLA dans le contexte d'un calcul TensorFlow simple:

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

Lorsqu'il est 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 afin qu'il calcule le résultat dans un seul lancement de noyau. Pour ce faire, il "fusionne" l'addition, la multiplication et la réduction dans un seul noyau GPU. De plus, cette opération fusionnée n'écrit pas en mémoire les valeurs intermédiaires produites par y*z et x+y*z. Au lieu de cela, elle "diffuse" les résultats de ces calculs intermédiaires directement vers leurs utilisateurs tout en les conservant entièrement 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. La suppression d'opérations mémoire est donc 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 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 utilise une sémantique must-compile: soit l'ensemble de la fonction est compilée avec XLA, soit une exception errors.InvalidArgumentError est générée. Actuellement, XLA ne peut pas compiler des fonctions dont les dimensions ne sont pas déduites, c'est-à-dire s'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)

Toutefois, 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 obtenir un exemple d'utilisation plus détaillé, ainsi qu'un tutoriel vidéo sur l'utilisation de 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 d'étape 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 aucune modification, une méthode simple consiste à activer le clustering automatique. Cette opération recherche automatiquement les clusters (sous-graphes connectés) dans les fonctions TensorFlow, 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

Le clustering automatique est actuellement optimisé pour les charges de travail GPU, mais vous pouvez également l'activer sur le processeur en utilisant l'option --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 sur le 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é. Les pièces jointes sont extrêmement utiles lorsque vous envoyez des rapports de bug XLA.

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

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

Vous pouvez également vider le graphe en visualisant la représentation vectorielle continue des clusters XLA à l'intérieur du graphe TensorFlow à l'aide de la commande 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 continue de clustering automatique utilisée. Si vous souhaitez les générer pour un programme TensorFlow s'exécutant avec le clustering automatique, lancez la commande suivante:

$ 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 run_hlo_module et en l'exécutant de manière itérative sur les programmes générés.

Complément d'informations

Interfaces XLA

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

  • JAX: transformations modulables 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