Pruna

Pruna est un cadre d’optimisation de modèles qui propose diverses méthodes d’optimisation – quantification, élagage, mise en cache, compilation – pour accélérer l’inférence et réduire l’utilisation de la mémoire. Un aperçu général des méthodes d’optimisation est présenté ci-dessous.

Technique	Description	Vitesse	Mémoire	Qualité
`batcher`	Regroupe plusieurs entrées pour qu'elles soient traitées simultanément, améliorant ainsi l'efficacité computationnelle et réduisant le temps de traitement.	✅	❌	➖
`cacher`	Stocke les résultats intermédiaires des calculs pour accélérer les opérations ultérieures.	✅	➖	➖
`compiler`	Optimise le modèle avec des instructions pour un matériel spécifique.	✅	➖	➖
`distiller`	Entraîne un modèle plus petit et plus simple pour imiter un modèle plus grand et plus complexe.	✅	✅	❌
`quantizer`	Réduit la précision des poids et des activations, ce qui diminue les besoins en mémoire.	✅	✅	❌
`pruner`	Supprime les connexions et les neurones moins importants ou redondants, ce qui aboutit à un réseau plus creux et plus efficace.	✅	✅	❌
`recoverer`	Restaure les performances d'un modèle après compression.	➖	➖	✅
`factorizer`	La factorisation regroupe plusieurs petites multiplications de matrices en une seule opération fusionnée de grande taille.	✅	➖	➖
`enhancer`	Améliore la sortie du modèle en appliquant des algorithmes de post-traitement tels que le débruitage ou le redimensionnement.	❌	-	✅

✅ (améliore), ➖ (approx. identique), ❌ (dégrade)

Explorez l’ensemble des méthodes d’optimisation dans la documentation Pruna.

Installation

Installez Pruna avec la commande suivante.

pip install pruna

Optimiser les modèles Diffusers

Une large gamme d'algorithmes d'optimisation est prise en charge pour les modèles Diffusers, comme indiqué ci-dessous.

Overview of the supported optimization algorithms for diffusers models

L’exemple ci-dessous optimise black-forest-labs/FLUX.1-dev à l’aide d’une combinaison d’algorithmes factorizer, compiler et cacher. Cette combinaison accélère l’inférence jusqu’à 4,2 fois et réduit la consommation mémoire GPU de pointe de 34,7 Go à 28,0 Go, tout en conservant une qualité de sortie pratiquement identique.

Consultez la documentation Pruna optimization pour en savoir plus sur les techniques d'optimisation utilisées dans cet exemple.

Optimization techniques used for FLUX.1-dev showing the combination of factorizer, compiler, and cacher algorithms

Commencez par définir un SmashConfig avec les algorithmes d'optimisation à utiliser. Pour optimiser le modèle, enveloppez le pipeline et le SmashConfig avec smash, puis utilisez le pipeline comme d'habitude pour l'inférence.

import torch
from diffusers import FluxPipeline

from pruna import PrunaModel, SmashConfig, smash

# load the model
# Try segmind/Segmind-Vega or black-forest-labs/FLUX.1-schnell with a small GPU memory
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

# define the configuration
smash_config = SmashConfig()
smash_config["factorizer"] = "qkv_diffusers"
smash_config["compiler"] = "torch_compile"
smash_config["torch_compile_target"] = "module_list"
smash_config["cacher"] = "fora"
smash_config["fora_interval"] = 2

# for the best results in terms of speed you can add these configs
# however they will increase your warmup time from 1.5 min to 10 min
# smash_config["torch_compile_mode"] = "max-autotune-no-cudagraphs"
# smash_config["quantizer"] = "torchao"
# smash_config["torchao_quant_type"] = "fp8dq"
# smash_config["torchao_excluded_modules"] = "norm+embedding"

# optimize the model
smashed_pipe = smash(pipe, smash_config)

# run the model
smashed_pipe("a knitted purple prune").images[0]

Après l'optimisation, nous pouvons partager et charger le modèle optimisé à l'aide du Hugging Face Hub.

# save the model
smashed_pipe.save_to_hub("<username>/FLUX.1-dev-smashed")

# load the model
smashed_pipe = PrunaModel.from_hub("<username>/FLUX.1-dev-smashed")

Évaluer et mesurer les performances des modèles Diffusers

Pruna fournit le EvaluationAgent pour évaluer la qualité de vos modèles optimisés.

Nous pouvons définir les métriques qui nous intéressent, telles que le temps total et le débit, ainsi que le jeu de données sur lequel évaluer. Nous pouvons définir un modèle et le passer à EvaluationAgent.

modèle optimisé

modèle autonome

Maintenant que vous avez vu comment optimiser et évaluer vos modèles, vous pouvez commencer à utiliser Pruna pour optimiser vos propres modèles. Heureusement, nous avons de nombreux exemples pour vous aider à démarrer.

Pour plus de détails sur le benchmarking de Flux, consultez le billet de blog Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)! et l'Espace InferBench.

Référence

Update on GitHub