AutoDeploy : vidages de graphes (AD_DUMP_GRAPHS_DIR)

Domaines d'application de cette skill

Ce fichier fait partie de trtllm-agent-toolkit. Les commandes et chemins tels que examples/auto_deploy/ et tensorrt_llm/ sont relatifs à un checkout de sources TensorRT-LLM, non au référentiel du plugin.

Quand utiliser cette skill

• Vous devez voir comment le graphe FX change après chaque transformation enregistrée exécutée.
• Vous vérifiez qu'un sous-graphe existe, qu'une fusion a correspondu, ou que des métadonnées / wrappers (getitem, view, reshape) sont apparus ou ont disparu entre les vidages.
• Vous appairez la sortie des journaux avec les fichiers de graphe sur disque lors du débogage d'AutoDeploy.

Skills connexes dans ce plugin

Variable d'environnement

Définissez :

export AD_DUMP_GRAPHS_DIR=/path/to/output/dir

Implémentation : GraphWriter.DUMP_GRAPHS_ENV == "AD_DUMP_GRAPHS_DIR" dans tensorrt_llm/_torch/auto_deploy/utils/graph_writer.py.

Si non définie, aucun fichier de graphe n'est écrit.

Quand les vidages sont produits

Après chaque application de transformation, BaseTransform appelle graph_writer.dump_graph(mod, t_name, self.config.stage.value) depuis tensorrt_llm/_torch/auto_deploy/transform/interface.py (immédiatement après _visualize_graph). Ainsi le vidage reflète le module après l'exécution de cette transformation.

Comportement du rang / processus

Depuis GraphWriter.dump_graph :

• Les vidages ne s'exécutent que si AD_DUMP_GRAPHS_DIR est défini.
• Si ADLogger.rank est défini et n'est pas 0, le vidage est ignoré (les processus non-rang-0 n'écrivent pas de fichiers).

Cycle de vie du répertoire

Au premier vidage sur le rang 0, GraphWriter supprime le répertoire cible s'il existe déjà, puis le recrée. Ne pointez pas AD_DUMP_GRAPHS_DIR vers un répertoire qui doit être conservé sans le copier d'abord.

Nommage et ordre des fichiers

Les fichiers sont nommés :

{NNN}_{<stage.value>}_{<transform_key>}.txt

• NNN est un compteur monotone de trois chiffres (001, 002, …) dans l'ordre d'exécution sur tous les vidages de ce processus.
• Le segment du milieu est la valeur config.stage de chaque transformation (le même enum/chaîne utilisé dans default.yaml sous le champ stage: de chaque transformation).
• Le dernier segment est la clé de registre de la transformation (transform_name passée à dump_graph).

Ainsi, un tri lexicographique par nom de fichier correspond à l'ordre du pipeline pour cette exécution.

Contenu des fichiers

Chaque fichier est en texte et commence par des en-têtes similaires à :

# Transform: <transform_key>
# Stage: <stage.value>
# GraphModules found: <count>

Ensuite, pour chaque torch.fx.GraphModule trouvé sous mod.named_modules() (y compris la racine), l'écrivain émet un titre de section et un listage de style SSA avec métadonnées de forme/dtype via dump_ssa_with_meta() du même module.

Utilisez ceci pour comparer les chaînes d'opérateurs, les consommateurs et les indices de forme/dtype node.meta entre fichiers consécutifs.

Exemple : capturer les vidages d'une construction de registre

Depuis la racine du clone TensorRT-LLM (adaptez le script et les drapeaux à votre workflow) :

AD_DUMP_GRAPHS_DIR=/tmp/ad-graphs \
  python examples/auto_deploy/build_and_run_ad.py --model <hf-model-id> --use-registry

Choisissez n'importe quel point d'entrée AutoDeploy que vous utilisez déjà ; l'exigence est seulement que le chemin de code exécute le pipeline de transformation avec AD_DUMP_GRAPHS_DIR défini dans l'environnement.

Journaux vs fichiers de vidage

Pendant l'exécution d'une transformation, la journalisation est patchée pour que les messages puissent être préfixés par [stage=<stage.value>, transform=<transform_key>] (voir with_transform_logging dans transform/interface.py). Les résumés de transformation enregistrent [SUMMARY] avec matches=<n> ou skipped / disabled (_log_transform_summary). Utilisez ces lignes ensemble avec les fichiers de vidage numérotés pour relier les comptages de correspondances aux formes de graphe avant et après une transformation spécifique.

Pièges

• Répertoire obsolète : Comme le répertoire de vidage est supprimé à la première utilisation, une deuxième exécution dans le même shell sans changer AD_DUMP_GRAPHS_DIR écrase la sortie antérieure.
• Aucun GraphModule : Si le module n'a pas d'enfants GraphModule, dump_graph retourne sans créer de nouveau fichier pour cette étape (voir le retour précoce dans graph_writer.py).
• Distribué : Seul le rang 0 écrit ; les autres rangs ignorent silencieusement.

Références sources

• tensorrt_llm/_torch/auto_deploy/utils/graph_writer.py — variable d'env, noms de fichiers, vidage SSA.
• tensorrt_llm/_torch/auto_deploy/transform/interface.py — site d'appel après chaque transformation ; décorateur de préfixe de journalisation.