Rendez vos ZeroGPU Spaces plus rapides avec la compilation AoT de PyTorch

TL;DR

• ZeroGPU permet d’allouer à la demande du matériel Nvidia H200 dans Hugging Face Spaces en forkant et libérant des processus GPU, réduisant le gaspillage des ressources en veille.
• La compilation ahead-of-time (AoT) de PyTorch accélère l’inférence de 1,3× à 1,8× sur des modèles comme Flux, Wan et LTX, avec environ 1,7× sur Flux.1-dev dans les Spaces ZeroGPU.
• Le flux AoT utilise spaces.aoti_capture pour capturer les entrées, torch.export.export pour générer le graphe compilé, puis spaces.aoti_compile et spaces.aoti_apply pour patcher le pipeline en vue de l’inférence.
• La quantification FP8 et les formes dynamiques offrent des gains supplémentaires; FA3 (Flash-Attention 3) peut être utilisé avec AoT pour des gains encore plus importants.
• Les utilisateurs Pro/Team/Enterprise bénéficient de 8× plus de quota ZeroGPU ; des exemples sont disponibles dans les espaces zerogpu-aoti; la compilation doit être effectuée dans un contexte @spaces.GPU en raison des dépendances matérielles.

Contexte et arrière-plan

Spaces est une plateforme Hugging Face permettant de publier des démos ML. Les applications typiques réservent un GPU pendant toute la durée du Space, même lorsque le trafic est minime. Lorsque .to(‘cuda’) est appelé, PyTorch initialise le pilote NVIDIA et maintient le processus sur CUDA, ce qui peut être inefficace avec un trafic irrégulier. ZeroGPU adopte une approche juste-à-temps: il fork le processus, initialiser CUDA pour la tâche GPU, exécuter la tâche et tuer le fork lorsque le GPU n’est plus nécessaire. Cela réduit le gaspillage et rend l’utilisation du GPU plus réactive. Actuellement, ZeroGPU alloue un MIG slice de H200 (profil 3g.71gb); des tailles MIG plus grandes (7g.141gb) sont prévues pour fin 2025. Les frameworks ML modernes, dont PyTorch et JAX, supportent des techniques de compilation qui optimisent la latence et le temps d’inférence via des transformations comme la fusion d’opérateurs. PyTorch 2.0 introduit deux interfaces majeures pour ce type de compilation: torch.compile, qui compile le modèle lors de la première exécution puis réutilise le résultat, et la compilation AoT. Sur ZeroGPU, torch.compile a du mal à réutiliser la compilation entre des processus fréquemment démarrés, s’appuyant sur le cache du système de fichiers, ce qui allonge les temps de démarrage. AoT exporte le modèle compilé une fois et le recharge instantanément dans de nouveaux processus. L’approche décrite utilise le transformeur Flux.1-dev comme exemple et se concentre sur la compilation du composant transformer, souvent le plus lourd dans les modèles génératifs. Cela s’aligne sur l’observation que le transformeur (ou le denoiser) porte la majeure partie du coût computationnel dans de nombreux modèles diffusion.

Qu’est-ce qui est compilé et pourquoi

AoT permet d’exporter un PyTorch program compilé qui représente le calcul et les paramètres du modèle, puis de le charger ultérieurement sans recompilation. Sur ZeroGPU, la compilation dépend fortement du matériel et nécessite des benchmarks, d’où l’enveloppement entier dans une fonction @spaces.GPU afin d’obtenir du code optimisé fiable.

Ce qui est nouveau

Le flux AoT dans ZeroGPU Spaces inclut:

• Capture d’entrées avec spaces.aoti_capture: utilisé comme contexte, il intercepte les appels du transformer ciblé (par exemple, pipe.transformer), empêche l’exécution et enregistre les args et kwargs pour l’export.
• Export vers PyTorch ExportedProgram: en utilisant les entrées capturées, on génère le graphe et les paramètres pour l’exportation.
• Compilation avec spaces.aoti_compile (wrapper de torch._inductor.aot_compile): produit un binaire AoT et gère le cycle de vie du binaire enregistré.
• Application du transformer compilé via spaces.aoti_apply: patch du forward du pipe.transformer pour utiliser le modèle compilé et suppression des anciens paramètres pour éviter OOM.
• Contexte GPU réel: comme la compilation dépend du matériel, tout doit se faire dans un @spaces.GPU.
• Résultat: un compiled_transformer prêt pour l’inférence; pour Flux.1-dev, environ 1,7× d’amélioration observée.

Remarques pratiques

• L’exemple de base utilise le modèle black-forest-labs/FLUX.1-dev et compile uniquement le transformer, car c’est souvent la partie la plus lourde.
• Le flux montre comment patcher le transformer sans perdre les attributs importants comme le dtype et la configuration.
• Il faut garder à l’esprit que l’installation et les pilotes CUDA dépendent fortement du matériel; la compilation AoT est donc une opération matérielle et doit être effectuée sur un GPU réel.

Quantification FP8 et formes dynamiques pour des gains supplémentaires

AoT peut être combiné avec la quantification pour des améliorations supplémentaires. Les schémas de quantification FP8 en post-formation dynamique offrent un bon compromis vitesse/qualité pour la génération d’images et de vidéos, à condition que la capacité CUDA soit d’au moins 9.0. Les GPUs H200 prennent en charge les schémas FP8 dans ce flux AoT, ouvrant la porte à des gains additionnels. Les API TorchAO facilitent l’intégration de la quantification FP8 dans le chemin AoT, apportant environ 1,2× de vitesse supplémentaire. Les formes dynamiques sont configurées via l’exportation et permettent de préciser quelles entrées varient, afin d’éviter des recompilations inutiles. Pour des cas où la variation est trop grande, on peut compiler un modèle par résolution et dispatcher le bon modèle au moment de l’exécution.

FA3 et accélération additionnelle

FA3 peut être utilisé dans les Spaces ZeroGPU en conjonction avec AoT pour accélérer davantage les charges liées à l’attention. La compilation et la construction de FA3 peuvent prendre du temps et dépendre du matériel, mais combinées à AoT, elles offrent des gains importants pour les charges d’attention.

Exemples et gains observés

Des ajustements relativement modestes dans le code peuvent transformer une démo en une démonstration très réactive. Dans le cas Flux.1-dev, le gain d’environ 1,7× a été observé; les gains globaux dépendent du modèle, des entrées et de l’usage de FP8 et des formes dynamiques. Les gains rapportés vont de 1,3× à 1,8× pour Flux, Wan et LTX, ce qui souligne l’utilité d’AoT dans le cadre ZeroGPU.

Multi-résolution et compilation par modèle

Pour les modèles où la résolution est critique, il peut être pertinent de compiler un modèle par résolution et de dispatcher le bon artefact à l’exécution, en conservant les paramètres partagés. Le fichier zerogpu-aoti-multi.py illustre un exemple minimal et Wan 2.2 Space démontre la mise en œuvre complète dans ce paradigme.

Pourquoi cela compte (impact pour les développeurs/entreprises)

• Démos plus rapides et plus réactives: AoT réduit les coûts de démarrage et améliore l’expérience utilisateur dans les demos hébergées sur Spaces.
• Utilisation efficiente du matériel ZeroGPU: évite les GPU inutilisés sur de longues périodes, alignant l’utilisation sur l’activité réelle.
• Combinaison d’options d’optimisation: FP8, formes dynamiques et FA3 offrent plusieurs leviers pour équilibrer vitesse et qualité.
• Capacité entreprise: les utilisateurs Pro/Team/Enterprise bénéficient de 8× de quota, facilitant l’expérimentation et le déploiement à l’échelle.
• Demos polyvalentes: le flux cible principalement les composants transformer, courants dans les modèles génératifs, rendant AoT largement applicable.

Détails techniques ou Implémentation

• Environnement: AoT nécessite un GPU, car les optimisations et benchmarks matériels sont spécifiques.
• Capture et export: utiliser spaces.aoti_capture pour capturer les entrées et exporter vers PyTorch ExportedProgram.
• Compilation: spaces.aoti_compile (wrapper de torch._inductor.aot_compile) produit le binaire AoT et gère le chargement lazily.
• Patch du pipeline: spaces.aoti_apply ajuste pipe.transformer.forward pour pointer vers le modèle compilé et libère les paramètres anciens pour éviter OOM.
• Contexte GPU: l’intégralité du flux AoT doit s’exécuter dans un @spaces.GPU pour tirer parti du matériel.
• Quantisation et formes dynamiques: configurer FP8 et définir des formes dynamiques selon le cas d’utilisation.
• Considérations par modèle: l’attention du transformer est un goulot d’étranglement typique; les modèles Wan et Flux servent d’exemples.
• Lectures de référence: Flux.1-dev et Wan sont présentés comme exemples; Wan 2.2 Space illustre le multival อย่าง multi-résolution. Le fichier zerogpu-aoti-multi.py montre une implémentation minimale.

Points clés (résumé)

• AoT permet de réutiliser immédiatement les modèles compilés, réduisant le temps de démarrage et améliorant l’expérience des démos sur Spaces.
• La quantification FP8 peut ajouter environ 1,2× de vitesse supplémentaire dans le flux AoT sur GPU H200.
• Les formes dynamiques offrent une flexibilité pour gérer des résolutions variables sans recompilation.
• En cas de variabilité élevée des entrées, compiler par résolution et dispatcher peut être avantageux.
• Des considérations matérielles signifient que les bénéfices de compilation se réalisent dans un environnement GPU; les outils de patch aident à la gestion de mémoire et d’état.

Références

• https://huggingface.co/blog/zerogpu-aoti