Étendre les exportateurs

torch.export trace le modèle en un graphe, les étapes ultérieures transforment ce graphe, et une étape finale le réduit ou l’émet pour le runtime cible. La plupart des modèles passent sans modification. Lorsqu’un motif PyTorch n’est pas géré par le backend, l’exporteur applique un petit contournement à cette étape plutôt que de modifier le modèle.

Ajoutez un contournement en écrivant une fonction et en l’enregistrant avec un décorateur. Chaque contournement doit être placé au niveau le plus bas qui permet de l’exprimer proprement.

Correctifs et modifications

Un contournement est soit un correctif, soit une solution. Les deux diffèrent par leur réversibilité.

Les deux sont enregistrées dans le registre de exporters/utils.py, et l’exportateur installe tout ce qui est enregistré pour son backend au bon moment.

Recourir à un correctif lorsque le problème est un bug de compilation propre à un seul backend : une traduction ONNX manquante, une particularité de validation ORT, ou une décomposition FX qui émet une opération morte. Le contournement reste dans l’exportateur, et le code de modélisation reste propre.

Lorsque les blocs de motifs sont exportés sur plusieurs backends, par exemple des boucles dépendantes des données, des caches avec état situés en dehors de Cache, ou une attention à boucle divisée écrite à la main, il est préférable de corriger le modèle. Une correction apportée une seule fois au modèle profite à tous les exportateurs.

Ajouter un patch

Supposons qu’une méthode de modèle fasse quelque chose que torch.export ne peut pas tracer. La méthode NllbMoeTop2Router._cast_classifier de NLLB-MoE convertit les poids du classifieur vers un autre dtype, ce qui n’est pas traçable. Remplacez-la par une opération nulle pour la durée de l’export.

Écrivez une fabrique qui prend la méthode d’origine et renvoie son remplacement, puis enregistrez la fabrique par rapport au chemin pointé de la méthode :

from transformers.exporters.utils import register_patch

@register_patch("dynamo", "transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier")
def _patch_classifier_cast(_original):
    # Replace the untraceable dtype cast with a no-op during export.
    return lambda self, *args, **kwargs: None

L'exportateur remplace la méthode avant le traçage et la restaure ensuite, de sorte que le correctif n'affecte que l'exportation. Quelques variantes :

• Passez des chemins supplémentaires pour partager une même usine entre plusieurs sites d’appel, par exemple @register_patch("dynamo", path_a, path_b).
• Patch une op torch en pointant le chemin vers elle, par exemple @register_patch("onnx", "torch.where"). L’usine reçoit l’op réelle en tant qu’argument, de sorte que le remplacement peut appeler celle-ci.
• Écrivez un correctif plutôt qu'un patch lorsque vous devez réécrire le graphe après le traçage. Le mécanisme est le même, une fonction décorée dans le registre correspondant.

Référence de stage

Les étiquettes de source de chaque exportateur marquent ses étapes par des blocs de commentaires # ── Stage N: … ──, de sorte que le fichier et cette référence correspondent. Consultez-les pour connaître les opérations et les classes exactes que chaque étape gère.

DynamoExporter

L'exportateur de base exécute une étape de patch et quatre utilitaires, dans l'ordre, dans DynamoExporter.export (voir exporter_dynamo.py).

1. Correctif de la signature avant : donne à model.forward une signature d’arguments plate afin que torch.export ne regroupe pas les entrées dans un seul tuple **kwargs. Il s’agit d’un élément interne et non d’un point d’extension.
2. Correctifs de modèle : remplacez les méthodes de modèle non traçables par des équivalents sûrs pour l’exportation pendant le traçage. Étendez avec @register_patch("dynamo", ...).
3. Enregistrement Pytree : enregistrez chaque Cache et ModelOutput afin que torch.export puisse les aplatir et les reconstruire (cela se produit généralement automatiquement). Ajoutez une branche à _flatten_to_context / _unflatten_from_context pour un type que la traversée des attributs ne peut pas atteindre.
4. Formes dynamiques : assignez Dim.AUTO à chaque tenseur et cachez la feuille lorsque dynamic=True. Remplacez par DynamoConfig.dynamic_shapes.
5. Nettoyage de l'état : réinitialiser les attributs de tenseurs qu'un modèle définit dans forward et que torch.export laisse en tant que tenseurs factices. Pour étendre, ajouter le nom de l'attribut à _STATEFUL_CACHE_ATTRS.

OnnxExporter

OnnxExporter ajoute cinq étapes autour de torch.onnx.export (voir exporter_onnx.py). Recherchez dans le fichier la liste complète des correctifs :

grep -nE "^def (_patch_|_fix_|_aten_)" src/transformers/exporters/exporter_onnx.py

1. Correctifs Torch : remplace les opérations torch que l’exportateur ONNX ne peut pas traduire telles quelles. Étendez avec @register_patch("onnx", ...).
2. Correctifs ONNX : relancer les corrections de nœuds après run_decompositions afin que les nœuds de garde de forme nouvellement introduits soient réparés avant l’abaissement. Utilise le même registre @register_patch("onnx", ...).
3. Corrections des nœuds FX : réécriture des nœuds du graphe que l’exporteur ONNX ne peut pas abaisser, tels que les opérations d’alias, les vues in-place et les assertions mortes. Étendez avec @register_fx_node_fix("onnx").
4. Traductions ONNX : fournir une conversion personnalisée pour une opération aten lorsque la conversion par défaut est absente ou bogueuse (par exemple aten.index_put ou aten._grouped_mm). Ajouter une fonction _aten_* à _ONNX_TRANSLATION_TABLE.
5. Corrections de l’IR ONNX : réécrire le programme ONNX terminé pour contourner les bogues d’ONNX Runtime (par exemple en forçant TopK(sorted=True)). Ajouter une fonction _fix_ir_* à _IR_FIXES.

ExecutorchExporter

ExecutorchExporter ajoute cinq étapes autour de to_edge_transform_and_lower et to_executorch, en commençant par la préparation du backend (voir exporter_executorch.py).

1. Préparation du backend : déplacer le modèle sur l'appareil et le type de données cibles et choisir son partitionneur (prepare_for_xnnpack, prepare_for_cuda). Ajouter un backend en enregistrant prepare_for_<name> dans _BACKEND_PREPARE.
2. Correctifs Torch : remplace les opérations torch que les backends ExecuTorch ne peuvent pas accepter, telles que split_copy, chunk et topk(k>dim). Étendez avec @register_patch("executorch", ...).
3. Patches ExecuTorch : remplace les composants internes d’ExecuTorch qui provoquent des plantages sur des graphes à formes dynamiques valides. Utilise le même registre @register_patch("executorch", ...).
4. Corrections du programme FX : réparez le programme exporté lorsque la correction nécessite un contexte de programme complet, comme l’élargissement des contraintes d’étendue ou le remplissage de métadonnées de placeholder manquantes. Étendez avec @register_fx_program_fix("executorch").
5. Corrections des nœuds FX : réécriture de nœuds individuels, par exemple la correspondance des opérations symboliques Python vers executorch_prim.* ou la réécriture de pow en une chaîne de mul. À étendre avec @register_fx_node_fix("executorch").

Contournements connus en amont

Certaines classes de modèles rencontrent des bugs confirmés dans l’optimiseur de graphes de onnxscript (plantage du repliement de constantes sur SplitToSequence, initialisateurs FPN supprimés). ONNX_DISABLE_OPTIMIZE désactive l’optimisation onnxscript pour ces modèles. Chaque entrée consigne le problème en amont à côté du nom du modèle. La liste est censée se réduire au fur et à mesure que les correctifs en amont sont intégrés ; une nouvelle entrée doit donc référencer un bug en amont spécifique plutôt que de désactiver l’optimisation arbitrairement.

EXPORT_SKIPS, exclut un petit nombre de classes de modèles du balayage d’exportation lorsque le modèle est fondamentalement non exportable en l’état (flux de contrôle dépendant des données qui ne peut pas être vectorisé, ou modules traités comme des arguments de forward). Chaque entrée contient une raison précisant la modification nécessaire côté modèle. Cette liste est également censée diminuer, et non augmenter.