Inférence LoRA rapide pour Flux avec Diffusers et PEFT

Les adaptateurs LoRA offrent une grande personnalisation pour des modèles de toutes formes et tailles. En matière de génération d’images, ils permettent d’apporter des styles différents, des personnages variés et bien plus encore. Parfois, ils peuvent aussi être utilisés pour diminuer la latence d’inférence. Ainsi, leur importance est primordiale, notamment pour la personnalisation et le fine-tuning des modèles. Dans cet article, nous utilisons le modèle Flux.1-Dev pour la génération de texte en image en raison de sa large popularité et adoption, et nous montrons comment optimiser sa vitesse d’inférence lors de l’utilisation de LoRA (~2,3x). Il compte plus de 30k adaptateurs entraînés avec lui (selon le Hub Hugging Face). Par conséquent, son importance pour la communauté est significative. Notez que même si nous montrons des gains avec Flux, notre conviction est que notre recette est suffisamment générique pour être appliquée à d’autres modèles. Si vous ne pouvez pas attendre le code, consultez le dépôt de code associé. Lors du service LoRA, il est courant d’effectuer du hot-swapping (échanger différents LoRA). Une LoRA modifie l’architecture du modèle de base. De plus, les LoRA peuvent être différentes entre elles – chacune peut avoir des rangs différents et viser des couches différentes pour l’adaptation. Pour tenir compte de ces propriétés dynamiques des LoRA, nous devons prendre les mesures nécessaires pour que les optimisations restent robustes. Par exemple, on peut appliquer torch.compile à un modèle chargé avec une LoRA particulière pour obtenir des gains d’inférence. Cependant, dès que l’on échange la LoRA contre une autre (avec une configuration potentiellement différente), on rencontre des problèmes de recompilation, ce qui ralentit l’inférence. On peut aussi fusionner les paramètres LoRA dans les paramètres du modèle de base, lancer la compilation et défusionner les LoRA lors du chargement de nouveaux. Cependant, cette approche rencontrera à nouveau des problèmes de recompilation à chaque inférence en raison de changements d’architecture potentiels. Notre recette d’optimisation prend en compte ces scénarios pour être aussi réaliste que possible. Parmi les composants clés, notez que la quantification FP8 est destructive mais offre le meilleur compromis vitesse/mémoire. Bien que la recette soit principalement testée sur des GPU NVIDIA, elle devrait fonctionner sur des GPU AMD aussi. Dans nos billets précédents, nous avons déjà discuté des bénéfices des trois premiers composants de notre recette d’optimisation. Leur application une par une ne nécessite que quelques lignes de code : Le processeur FA3 provient d’ici. Les problèmes apparaissent lorsque nous essayons d’échanger des LoRAs dans un transformateur de diffusion compilé (pipe.transformer) sans déclencher de recompilation. En général, charger et décharger des LoRAs nécessiterait une recompilation, ce qui annulerait les gains de vitesse obtenus par la compilation. Heureusement, il existe une solution pour éviter la recompilation. En passant hotswap=True, diffusers laissera l’architecture du modèle inchangée et n’échangera que les poids de l’adaptateur LoRA, ce qui ne nécessite pas de recompilation. (Pour rappel, le premier appel à pipe sera lent car torch.compile est un compilateur JIT. Cependant, les appels suivants seront fortement plus rapides.) Cela permet généralement d’échanger les LoRA sans recompilation, mais il existe des limitations : pour plus d’informations sur le hot-swapping dans Diffusers et ses limitations, consultez la section hot-swapping de la documentation. Les bénéfices de ce flux deviennent évidents lorsque l’on observe la latence d’inférence sans compilation avec hot-swapping. Points clés : Cette recette d’optimisation suppose l’accès à un GPU puissant comme le H100. Mais que faire si l’on dispose uniquement d’un GPU grand public comme le RTX 4090 ? Découvrons-le. Flux.1-Dev (sans LoRA) utilisant le type de données BF16 nécessite ~33 Go de mémoire pour fonctionner. Selon la taille du module LoRA et sans optimisation, ce footprint peut augmenter encore. Dans ce contexte, la RTX 4090 est notre matériel de test. Premièrement, pour exécuter Flux.1-Dev de bout en bout, on peut appliquer un offloading CPU où les composants non nécessaires à l’exécution courante sont déportés sur le CPU afin de libérer plus de mémoire sur l’accélérateur. Cela permet d’exécuter l’intégralité du pipeline en ~22 Go en 35,403 secondes sur une RTX 4090. Activer la compilation peut réduire la latence à 31,205 secondes (gain ~1,12x). En termes de code, cela représente quelques lignes seulement : remarquez que nous n’avons pas appliqué la quantisation FP8 ici car elle n’est pas prise en charge avec l’offload CPU et la compilation (fil d’assistance). Par conséquent, se limiter à la quantisation FP8 du Flux Transformer ne suffit pas à atténuer le problème de mémoire épuisée. Pour profiter de FP8 sans offload CPU, il faut quantifier le T5. Pour Flux.1-Dev, si nous quantifions en plus l’encodeur de texte T5, nous devrions pouvoir charger et exécuter l’ensemble du pipeline en 24 Go. Ci-dessous, une comparaison des résultats avec et sans quantification de l’encodeur T5 (NF4 de bitsandbytes). Comme on peut le voir, quantifier l’encodeur T5 n’entraîne pas une perte de qualité notable. Combiner l’encodeur T5 quantifié et le Flux Transformer quantifié avec FP8 et torch.compile donne des résultats raisonnables – 9,668 secondes contre 32,27 secondes (gain d’environ ~3,3x) sans perte perceptible de qualité. Il est possible de générer des images avec 24 Go de VRAM même sans quantifier l’encodeur T5, mais cela rendrait le pipeline légèrement plus complexe. Nous disposons désormais d’un moyen d’exécuter l’intégralité du pipeline Flux.1-Dev avec quantisation FP8 sur une RTX 4090. Nous pouvons appliquer la recette d’optimisation préétablie pour optimiser l’inférence LoRA sur ce même matériel. Comme FA3 n’est pas pris en charge sur la RTX 4090, nous restons sur la recette suivante avec quantisation du T5 ajoutée récemment : Dans le tableau ci-dessous, nous montrons les latences d’inférence pour différentes combinaisons des composants appliqués. Remarques rapides : Pour activer le hot-swapping sans déclencher de recompilation, deux obstacles doivent être surmontés. Premièrement, le facteur d’échelle LoRA doit être converti en tenseurs PyTorch à partir de nombres à virgule flottante, ce qui se fait assez simplement. Deuxièmement, la forme des poids LoRA doit être remplie jusqu’à la plus grande forme requise. Ainsi, les données dans les poids peuvent être remplacées sans réaffecter tout l’attribut. C’est pourquoi l’argument max_rank est crucial. En remplissant avec des zéros, les résultats restent inchangés, bien que le calcul puisse être ralenti selon la taille du remplissage. Comme aucun nouvel attribut LoRA n’est ajouté, il faut que chaque LoRA après la première cible les mêmes couches, ou un sous-ensemble des couches, que la première LoRA. Choisissez donc l’ordre de chargement avec soin. Si les LoRAs ciblent des couches disjointes, il est possible de créer une LoRA fictive qui cible l’union de toutes les couches cibles. Pour voir les détails de cette implémentation, consultez le fichier hotswap.py dans PEFT. Cet article décrit une recette d’optimisation pour une inférence LoRA rapide avec Flux, démontrant des gains significatifs de vitesse. Notre approche combine Flash Attention 3, torch.compile et quantisation FP8 tout en garantissant la possibilité de hot-swapping sans recompilation. Sur des GPUs haut de gamme comme le H100, cette configuration optimise environ 2,23x plus vite que le baseline. Pour les GPUs grand public, en particulier le RTX 4090, nous avons géré les contraintes mémoire en introduisant la quantisation NF4 du texte T5 et en utilisant une compilation régionale. Cette recette complète a abouti à un gain substantiel d’environ 2,04x, rendant l’inférence LoRA sur Flux viable et performante même avec une VRAM limitée. L’idée clé est que, en gérant soigneusement compilation et quantisation, les bénéfices de LoRA peuvent être pleinement réalisés sur différentes configurations matérielles. Nous espérons que ces recettes vous inspireront pour optimiser vos cas d’utilisation basés sur LoRA et profiter d’une inférence rapide. Ci-dessous, una liste des ressources importantes citées tout au long de cet article: Plus d’articles de notre Blog Par linoyts 2 janvier 2024 • 70 Par derekl35 19 juin 2025 • 83 J’ai aimé votre article ! Puis-je demander quelle version de PyTorch vous avez utilisée pour les résultats de test ? PyTorch nightly. A boy with baloon · Inscrivez-vous ou connectez-vous pour commenter