# Pruna

[Pruna](https://github.com/PrunaAI/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](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms).

## Installation

Installez Pruna avec la commande suivante.

```bash
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.

    

L’exemple ci-dessous optimise [black-forest-labs/FLUX.1-dev](https://huggingface.co/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.

> [!TIP]
> Consultez la documentation [Pruna optimization](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html) pour en savoir plus sur les techniques d'optimisation utilisées dans cet exemple.

    

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.

```python
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.

```python
# 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](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html) 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`.

Nous pouvons charger et évaluer un modèle optimisé en utilisant le `EvaluationAgent` et en le transmettant à la `Task`.

```python
import torch
from diffusers import FluxPipeline

from pruna import PrunaModel
from pruna.data.pruna_datamodule import PrunaDataModule
from pruna.evaluation.evaluation_agent import EvaluationAgent
from pruna.evaluation.metrics import (
    ThroughputMetric,
    TorchMetricWrapper,
    TotalTimeMetric,
)
from pruna.evaluation.task import Task

# define the device
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
smashed_pipe = PrunaModel.from_hub("PrunaAI/FLUX.1-dev-smashed")

# Define the metrics
metrics = [
    TotalTimeMetric(n_iterations=20, n_warmup_iterations=5),
    ThroughputMetric(n_iterations=20, n_warmup_iterations=5),
    TorchMetricWrapper("clip"),
]

# Define the datamodule
datamodule = PrunaDataModule.from_string("LAION256")
datamodule.limit_datasets(10)

# Define the task and evaluation agent
task = Task(metrics, datamodule=datamodule, device=device)
eval_agent = EvaluationAgent(task)

# Evaluate smashed model and offload it to CPU
smashed_pipe.move_to_device(device)
smashed_pipe_results = eval_agent.evaluate(smashed_pipe)
smashed_pipe.move_to_device("cpu")
```

Plutôt que de comparer le modèle optimisé au modèle de base, vous pouvez également évaluer le modèle `diffusers` autonome. Cela est utile si vous souhaitez évaluer les performances du modèle sans l’optimisation. Nous pouvons le faire en utilisant l’enveloppe `PrunaModel` et en exécutant le `EvaluationAgent` dessus.

```python
import torch
from diffusers import FluxPipeline

from pruna import PrunaModel

# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16
).to("cpu")
wrapped_pipe = PrunaModel(model=pipe)
```

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.

> [!TIP]
> 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)!](https://huggingface.co/blog/PrunaAI/flux-fastest-image-generation-endpoint) et l'Espace [InferBench](https://huggingface.co/spaces/PrunaAI/InferBench).

## Référence

- [Pruna](https://github.com/PrunaAI/pruna)
- [Optimisation Pruna](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms)
- [Évaluation Pruna](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)
- [Tutoriels Pruna](https://docs.pruna.ai/en/stable/docs_pruna/tutorials/index.html)