Vider les calculs HLO

Un dump HLO est une représentation textuelle des modules HLO à différentes étapes du calcul. Il est utile pour le débogage et vous devez souvent l'inclure dans les rapports de bug. Il s'agit généralement d'un fichier texte lisible qui liste les instructions HLO et leurs propriétés. Parfois, les modules HLO sont vidés comme suit :

  • HloProto : fichiers de tampon de protocole, qui sont un format plus structuré et lisible par une machine.
  • HloSnapshot : module HLO et ses entrées. Pour relire les HLO, vous avez parfois besoin des entrées réelles fournies à un calcul donné plutôt que de données aléatoires.

Vous pouvez utiliser des indicateurs XLA pour spécifier et obtenir des dumps. Dans la plupart des cas, vous pouvez le définir avec une variable d'environnement. JAX propose également une méthode programmatique pour imprimer le dump HLO.

Exécution locale

Utiliser des variables d'environnement

Vous pouvez définir la variable d'environnement XLA_FLAGS avec les indicateurs nécessaires pour obtenir des dumps. Cela fonctionne pour JAX, TensorFlow et PyTorch/XLA.

Pour vider les modules HLO et d'autres informations de débogage dans un répertoire spécifique, exécutez votre programme avec l'indicateur --xla_dump_to :

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Par exemple, vous pouvez utiliser /tmp ou /tmp/xladump comme chemins d'accès.

Par défaut, cela génère des modules HLO sous forme de texte, au tout début et à la toute fin du pipeline d'optimisation.

Vous pouvez également spécifier explicitement le format :

  1. Dumps de texte
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. Fichiers .proto HLO
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. Instantanés HLO
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. Rendu du graphique avec le serveur graphviz (ne fonctionne bien que pour les petits graphiques)
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. Rendu du graphique dans un fichier HTML (ne fonctionne bien que pour les petits graphiques)
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

Pour les graphiques plus volumineux, vous pouvez utiliser interactive_graphviz pour visualiser des parties du graphique.

Afficher les passes intermédiaires spécifiques

En plus des HLO pré-optimisés / optimisés finaux standards, vous pouvez également vider l'état des HLO après un pass de compilation spécifique.

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

Les modules HLO seront vidés pour les passes dont les noms correspondent à l'expression régulière (regex). Par exemple, vous pouvez observer les HLO résultant des passes liées au partitionnement SPMD avec :

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

Pour vider le résultat après chaque passe XLA (cela générera de nombreux fichiers), vous pouvez définir :

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

Options spécifiques à JAX

De manière programmatique dans JAX

Au lieu de transmettre des indicateurs ou des variables d'environnement, vous pouvez également vider le HLO de manière programmatique à l'aide des API lower et compile de JAX.

Récupérez localement le HLO abaissé d'origine non optimisé avec :

jax.jit(f).lower(*args).as_text('hlo')

Pour effectuer un dump dans des fichiers lors des passes de compilation HLO, spécifiez :

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

Vider les jaxprs

Les jaxpr sont la représentation intermédiaire de JAX pour les traces de programme. Pour ce faire, définissez les variables d'environnement suivantes :

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Pour en savoir plus, consultez la documentation JAX sur l'exportation et la sérialisation des calculs mis en scène : débogage.

Google Colab

Variables d'environnement

Dans la première cellule exécutée de votre notebook (car les variables d'environnement et les indicateurs de ligne de commande ne sont généralement traités qu'une seule fois, par exemple au moment de l'importation du module ou de l'initialisation du backend XLA), ajoutez le XLA_FLAGS ci-dessus avec os.environ, par exemple :

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

Cela permet de transférer le calcul vers DIRECTORY_PATH, par exemple /tmp. Dans Colab, accédez au navigateur "Fichiers" dans la barre latérale de gauche pour afficher ce répertoire et y accéder.

Vous pouvez utiliser tous les indicateurs mentionnés dans la section "Exécution locale".

Options spécifiques à JAX

Comme pour l'exécution locale, vous pouvez imprimer directement le HLO pré-optimisé d'un calcul pour une introspection interactive en direct :

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

Vous pouvez également imprimer directement le HLO optimisé d'un calcul :

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

Vider tous les calculs/une partie des calculs

Si vous souhaitez tout voir dans un dump, y compris toutes les petites compilations, définissez la variable d'environnement JAX :

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mosaïque

Mosaic est un compilateur pour le backend Pallas TPU et le backend Pallas GPU expérimental. Pour vider le calcul du mosaïque, définissez l'indicateur suivant :

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Vous pouvez également définir les arguments d'initialisation du TPU en tant que variable d'environnement :

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

Pour en savoir plus, consultez la documentation JAX sur Pallas et Mosaic.

Autres contenus avec HLO Dumps

Trouver le bon calcul

En général, de nombreux calculs sont abandonnés. Les fichiers dumpés sont nommés de manière explicite avec le "nom de calcul" JAX, TensorFlow ou PyTorch/XLA indiqué dans les journaux, ce qui permet d'identifier facilement les fichiers HLO pertinents. Exemple :

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

Sinon, vous pouvez utiliser ripgrep pour identifier rapidement le module contenant des symboles ou des calculs spécifiques.

Conseil : Incluez les trois fichiers de vidage avant/après/d'attribution de tampon qui vous intéressent dans vos rapports de bug.

Conversion HLO

Un outil appelé hlo-opt qui peut traduire les formats HLOProto et texte. Cette fonctionnalité est utile lorsque vous disposez d'un format, mais que vous avez besoin de l'autre pour le débogage.

Apprenez à l'utiliser : Documentation sur l'outil XLA : hlo-opt.

Revoir

Vous pouvez exécuter (relire) les calculs dumpés sur un backend XLA spécifié avec de fausses données ou des instantanés d'entrée. Il s'agit d'un moyen pratique de reproduire, d'itérer et de déboguer les problèmes dans XLA.

Les commandes suivantes utilisent de fausses données. Si vous avez enregistré des instantanés HLO, vous pouvez les transmettre à la place. Les données de l'instantané seront alors utilisées. Pour continuer à utiliser de fausses données lors de l'exécution de l'instantané, transmettez l'option --force_fake_data.

Backend du CPU :

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

Backend GPU :

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt