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

TL;DR

• ZeroGPU Spaces peut désormais tirer parti de la compilation ahead-of-time (AoT) de PyTorch pour réduire la latence et les temps de démarrage à froid des tâches GPU.
• AoT permet d’exporter un modèle compilé une fois et de le recharger instantanément dans de nouveaux processus, offrant des démonstrations plus réactives sur Hugging Face Spaces avec le matériel Nvidia H200.
• Des gains réels d’environ 1,3× à 1,8× ont été observés sur des modèles tels que Flux.1-dev, Wan et LTX; la quantification FP8 peut ajouter environ 1,2× de vitesse supplémentaire.
• Le flux de travail s’appuie sur spaces.aoti_capture et spaces.aoti_compile, plus spaces.aoti_apply pour patcher le pipeline en toute sécurité sans conserver de gros paramètres du modèle en mémoire.
• MIG slices actuels sur ZeroGPU: 3g.71gb; une tranche complète 7g.141gb est attendue fin 2025.
• ZeroGPU peut être combiné avec Flash-Attention 3 (FA3) pour accélérer davantage les charges de travail. Pour plus de détails, consultez le post Hugging Face d’origine. Source

Contexte et antécédents

Spaces est une plateforme Hugging Face permettant aux praticiens ML de publier des démos d’applications avec des GPU puissants. ZeroGPU évite de réserver un GPU en permanence pour chaque Space, en allouant la ressource GPU uniquement en cas d’activité, ce qui optimise les ressources lorsque le trafic est clairsemé et irrégulier. L’approche traditionnelle oblige à initialiser le pilote NVIDIA lors de l’appel .to(‘cuda’), ce qui fixe le processus CUDA de manière permanente. ZeroGPU adopte une approche de démarrage à la demande: il bifurque le processus, configure CUDA pour la tâche, exécute le travail et libère le GPU ensuite. Cette approche repose sur le package spaces et sur le décorateur GPU. Mais AoT apporte une optimisation importante: elle permet d’exporter un programme compilé et de le réutiliser rapidement sans recompilation, ce qui réduit les surcoûts et les temps de démarrage. Référence à l’article d’origine sur AoT par Hugging Face. Source

Ce qui est nouveau

L’objectif est de compiler à l’avance le composant transformer (la partie denoiser dans de nombreux modèles génératifs), car c’est le cœur du coût de calcul. Le flux intègre plusieurs éléments du package spaces:

• aoti_capture: utilisé comme gestionnaire de contexte pour intercepter les appels du transformer et enregistrer les entrées simulées.
• torch.export.export: exporte le transformer en PyTorch ExportedProgram, qui est un graphe de calcul avec les paramètres.
• spaces.aoti_compile: wrapper autour de torch._inductor.aot_compile qui gère l’enregistrement et le chargement paresseux du modèle compilé.
• spaces.aoti_apply: patch sûr du pipeline en remplaçant forward par le modèle compilé et en nettoyant les paramètres pour éviter les fuites mémoire.
• Encadrement par @spaces.GPU: tout doit être exécuté dans ce cadre pour bénéficier du matériel et du tuning spécifiques à AoT.
• Compatibilité: remplacer directement par compiled_transformer peut supprimer des attributs importants; spaces.aoti_apply assure une intégration propre. Concrètement, à partir de l’exemple ZeroGPU de base et en utilisant le transformeur Flux.1-dev, on observe environ 1,7× d’accélération avec AoT sur ZeroGPU. Cela démontre qu’un seul passage de compilation AoT peut apporter des gains significatifs pour les démonstrations dans Spaces. La quantification FP8 peut être associée à AoT pour des gains supplémentaires. FP8 nécessite une compute capability CUDA d’au moins 9.0, ce qui est satisfait par les GPUs H200 dans ZeroGPU. TorchAO rend ce flux possible et permet d’explorer la quantification et les formes dynamiques lors de l’exportation (transformer_dynamic_shapes). Pour Flux.1-Dev, les entrées d’exemple et les plages dynamiques sont déterminées par l’analyse des shapes de hidden_states sous des résolutions variées. Dans les cas où le dynamisme est très important (comme Wan pour la génération vidéo), on peut compiler un modèle par résolution et dispatcher le bon artefact au moment de l’exécution. FA3 (Flash-Attention 3) est compatible avec AoT, offrant davantage d’accélérations pour les calculs d’attention; l’installation et l’optimisation de FA3 peuvent prendre du temps et dépendent du hardware. Des démonstrations complètes existent dans zerogpu-aoti-multi.py et dans Wan 2.2 Space. Ensemble, AoT, quantification FP8 et formes dynamiques offrent un cadre pratique pour accélérer les démos sur Spaces avec ZeroGPU.

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

AoT avec ZeroGPU Spaces permet des démonstrations plus rapides sans nécessiter des GPU continuellement allumés. La latence réduite au démarrage bénéficie grandement aux flux de démonstration et améliore l’expérience utilisateur final. Sur le plan architectural, compiler puis réutiliser le artefact AoT s’aligne avec l’objectif de ZeroGPU d’utilisation efficace des ressources GPU à la demande. Les entreprises et équipes utilisant Spaces obtiennent plusieurs avantages:

• Démos plus rapides et cycles d’itération plus courts.
• Meilleure utilisation des ressources en évitant les initialisations CUDA répétées.
• Potentiel de throughput accru grâce à la quantification FP8 et aux formes dynamiques.
• Option de gestion par résolution lorsque le dynamisme est élevé, en conservant les paramètres partagés. Les informations techniques viennent du post Hugging Face original; consultez la source pour les détails. Source

Détails techniques ou Mise en œuvre

La mise en œuvre décrit un flux concret dans un Spaces ZeroGPU:

• Préparer des entrées d’exemple via spaces.aoti_capture pour représenter les tenseurs d’inférence réels.
• Exporter le transformer en PyTorch ExportedProgram avec torch.export.export.
• Compiler via spaces.aoti_compile, qui encapsule aot_compile et gère le chargement paresseux du modèle compilé.
• Appliquer le patch via spaces.aoti_apply pour remplacer forward par le modèle compilé et nettoyer les paramètres.
• Encapsuler le tout dans une fonction @spaces.GPU afin d’assurer l’accès au matériel réel et le tuning matériel.
• L’inférence via AoT: le transformer compilé peut être réutilisé sans réexportation, réduisant le coût total par exécution. Knobs pratiques:
• Quantification FP8 avec TorchAO: accélère les générations d’images/vidéos en complément d’AoT, sous réserve de la compatibilité matérielle.
• Formes dynamiques: exportez avec transformer_dynamic_shapes pour traiter facilement différentes résolutions d’entrée.
• Wan et résolutions multiples: lorsque le dynamisme est élevé, compiler par résolution et dispatcher le bon artefact à l’exécution.
• FA3: compatible avec AoT pour accélérer l’attention; la construction et le tuning peuvent prendre du temps et dépendre du hardware. Notes matérielles et capacités:
• MIG slices: ZeroGPU dispose actuellement d’un MIG slice de 3g.71gb; un slice complet 7g.141gb est prévu pour fin 2025.
• Compatibilité: AoT nécessite un environnement CUDA réel et ne peut être entièrement simulé hors du frame @spaces.GPU. Pour ceux qui veulent commencer, les exemples Flux.1-Dev, Wan et LTX illustrent les gains pratiques avec AoT. Le post Hugging Face d’origine offre une vue détaillée du processus. Source

Points clés

• AoT permet la réutilisation immédiate du modèle compilé entre processus ZeroGPU, réduisant la latence de démarrage.
• L’intégration avec aoti_capture, aoti_compile et aoti_apply rend l’amélioration accessible même pour des demos existantes dans Spaces.
• FP8 et formes dynamiques apportent des gains additionnels et élargissent les scénarios d’utilisation.
• Wan montre une approche par résolution lorsque le dynamisme est élevé, tout en réutilisant les paramètres partagés.
• FA3 est compatible et peut offrir des accélérations supplémentaires sur les calculs d’attention.

Références

• Article Hugging Face sur ZeroGPU AoT: https://huggingface.co/blog/zerogpu-aoti