Accélérez ZeroGPU Spaces avec la compilation ahead-of-time (AoT) de PyTorch

Aperçu

ZeroGPU permet d’utiliser du matériel Nvidia H200 puissant dans Hugging Face Spaces sans réserver la GPU en permanence. Bien que c’est efficace pour les démos, l’optimisation complète du stack CUDA n’est pas toujours atteinte. La compilation ahead-of-time (AoT) de PyTorch résout ce problème en exportant un modèle compilé une fois et en le rechargeant instantanément dans de nouveaux processus, ce qui rend les démos plus réactives et l’expérience plus fluide. Concrètement, AoT offre des gains typiques de 1,3× à 1,8× sur des modèles tels que Flux.1-dev, Wan et LTX dans ZeroGPU Spaces. AoT est particulièrement utile pour le modèle d’exécution de ZeroGPU, basé sur des processus qui se montent rapidement et s’éteignent après usage. Plutôt que de compiler à chaque exécution (ce qui est coûteux lorsque les processus sont créés pour chaque tâche GPU), AoT précompile et réutilise le graphe optimisé. L’article montre comment intégrer AoT dans les ZeroGPU Spaces, avec des astuces avancées comme la quantification FP8 et les formes dynamiques, et fournit des démos prêtes à l’emploi. ZeroGPU exploite les sections MIG de l’H200 et des capacités futures, avec des variantes comme 3g.71gb actuellement disponibles et 7g.141gb prévues pour fin 2025. L’approche vise à réduire les coûts mémoire et à allouer la GPU uniquement lorsque nécessaire.

Caractéristiques clés

AoT exporte et recharge: compiler un modèle une fois et réutiliser le programme compilé dans différents processus et tâches, réduction de la latence à l’exécution.
Intégration Spaces: utilisez le package spaces avec le décorateur GPU pour activer le flux AoT dans ZeroGPU.
Outils AoT: spaces.aoti_capture pour capturer des entrées d’exemple, spaces.aoti_compile pour générer le binaire AoT, et spaces.aoti_apply pour patcher le pipeline en toute sécurité.
Noyau PyTorch: AoT s’appuie sur torch._inductor.aot_compile et nécessite un PyTorch ExportedProgram compilé.
Quantification: la quantification FP8 en post-entraînement peut être associée à AoT pour des gains supplémentaires (nécessite une compute capability 9.0+).
Formes dynamiques: l’export peut être configuré pour traiter certaines entrées comme dynamiques, permettant de gérer différentes résolutions d’image/vidéo.
Modèle par résolution: pour la génération vidéo très dynamique (Wan), vous pouvez compiler un modèle par résolution tout en partageant les paramètres.
Compatibilité FA3: Flash Attention 3 peut être utilisé dans les flux AoT pour des améliorations de performance supplémentaires.
Matériel: ZeroGPU fonctionne sur des H200, avec des MIG slices tels que 3g.71gb; d’autres tailles sont prévues.
Démonstrations de vitesse: un exemple montre ~1,7× d’amélioration sur Flux.1-dev; 1,3×–1,8× selon le modèle.

L’article lie aussi à TorchAO et tutoriels PyTorch AoTInductor pour plus de contexte.

Cas d’usage courants

Démos en direct et showcases sur Hugging Face Spaces, lorsque l’activité utilisateur est sporadique et que l’utilisation de la GPU est coûteuse.
Tâches de génération d’images et de vidéos (familles Flux et Wan) où la latence d’inférence est critique pour des démos interactives.
Scénarios nécessitant une initialisation rapide des modèles compilés entre différents Spaces/tasks, évitant les surcoûts de compilation par tâche.
Flux et flux quantifiés FP8 et formes dynamiques pour maximiser le throughput sur le matériel H200.
Génération vidéo multi-résolution avec dispatch du bon modèle compilé tout en partageant les paramètres.

Installation et configuration

L’article décrit comment activer AoT au sein de ZeroGPU Spaces en utilisant le package spaces et les outils Torch Inductor. Le workflow implique d’intercepter les entrées, d’exporter le module vers un PyTorch ExportedProgram, de le compiler avec AoT, puis de remplacer le transformer par la version compilée dans le pipeline. Pour la compilation, un GPU réel est nécessaire, car l’optimisation dépend du matériel. Étapes clés (conceptuelles) :

Envelopper le calcul dans une fonction @spaces.GPU pour assurer la gestion du dispositif et activer le cycle de vie JIT de ZeroGPU.
Utiliser spaces.aoti_capture pour intercepter les entrées d’exemple du transformer.
Exporter le transformer avec torch.export.export.
Compiler l’objet exporté avec spaces.aoti_compile, qui s’appuie sur torch._inductor.aot_compile et gère le chargement/stockage.
Appliquer le patch via spaces.aoti_apply pour remplacer forward du transformer par le modèle compilé et libérer les paramètres obsolètes.
Exécuter l’inférence dans le contexte Spaces.GPU pour garantir la compatibilité et tirer parti du graphe AoT compilé. Remarque: l’article souligne que les commandes exactes dépendent de l’environnement et qu’un GPU réel est nécessaire pendant la compilation.

# Commandes exactes non fournies dans la source. Voir l’article pour le flux complet.
# Exemplaires placeholders (non tirés textuellement de la source) :
# 1. from spaces import GPU, aoti_capture, aoti_compile, aoti_apply
# 2. @GPU
# def run(pipe):
# with aoti_capture(pipe.transformer) as call:
# pass
# compiled = aoti_compile(call)
# aoti_apply(pipe, compiled)
# # exécuter l’inférence ici

Démarrage rapide (aperçu minimal et exécutable)

Identifiez le composant transformer lourd dans votre pipeline de diffusion/génération (par exemple dans Flux.1-dev).
Encapsuler le code de démonstration dans une fonction Spaces.GPU pour activer le comportement JIT de ZeroGPU.
Capturer les entrées avec aoti_capture, exporter avec torch.export.export, compiler avec aoti_compile et appliquer avec aoti_apply.
Exécuter une inférence de test pour vérifier les gains et s’assurer que le chemin compilé est utilisé. L’article indique qu’un transformer compilé a donné environ 1,7× d’amélioration dans Flux.1-dev et que la quantification FP8 peut ajouter ~1,2× supplémentaires lorsque c’est pertinent.

Avantages et inconvénients

Avantages
Réduction significative de la latence de démarrage/ compilation pour les ZeroGPU Spaces.
Flux clair entre exportation, compilation et exécution.
Possibilité de combiner AoT avec quantification FP8 et formes dynamiques pour des gains supplémentaires.
Compatibilité FA3 pour des accélérations supplémentaires dans les modèles de génération.
Compatible avec l’allocation à la demande de ZeroGPU, maximisant l’efficacité des ressources.
Inconvénients
Nécessite un GPU réel pendant la compilation (tuning matériel-dépendant).
Certains modèles peuvent nécessiter des variantes compilées par résolution pour des résultats optimaux.
La configuration implique plusieurs étapes (capture, export, compilation, apply) qui doivent être synchronisées entre les Spaces.

Alternatives (comparaison rapide)

| Approche | Avantages | Inconvénients |---|---|---| | AoT avec les outils spaces.aoti_* | Recharge rapide entre processus; évite la compilation par exécution; support des formes dynamiques et FP8 | Nécessite un GPU pendant la compilation; intégration multi-API; nécessite une séparation des composants |PyTorch torch.compile (JIT) | chemin plus simple en environnements standard; compilation lors de la première exécution et cache sur le filesystem | Sur ZeroGPU, chaque tâche démarre un nouveau processus, ce qui limite l’utilité du cache FS; démarrage à froid persiste |Inférence non AoT | Simple; pas de complexité d’export/compile | Plus de latence à l’initialisation sur ZeroGPU |