Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /v0.43.0 /zh /optimization /pruna.md

mishig

2 days ago

preview code

download

raw

7.31 kB

# Pruna

[Pruna](https://github.com/PrunaAI/pruna) 是一个模型优化框架，提供多种优化方法——量化、剪枝、缓存、编译——以加速推理并降低内存占用。下图展示了这些优化方法的总体概览。

| 技术    | 描述                                                                                   | 速度 | 内存 | 质量 |
|--------------|-----------------------------------------------------------------------------------------------|:-----:|:------:|:-------:|
| `batcher`    | 将多个输入组合在一起以同时处理，从而提高计算效率并缩短处理时间。 | ✅    | ❌     | ➖      |
| `cacher`     | 存储计算的中间结果，以加速后续操作。               | ✅    | ➖     | ➖      |
| `compiler`   | 针对特定硬件优化模型。                                 | ✅    | ➖     | ➖      |
| `distiller`  | 训练一个更小、更简单的模型，以模拟一个更大、更复杂的模型。                       | ✅    | ✅     | ❌      |
| `quantizer`  | 降低权重和激活值的精度，从而减少内存需求。              | ✅    | ✅     | ❌      |
| `pruner`     | 移除不太重要或冗余的连接和神经元，从而得到一个更稀疏、更高效的网络。 | ✅    | ✅     | ❌      |
| `recoverer`  | 恢复模型压缩后的性能。                                       | ➖    | ➖     | ✅      |
| `factorizer` | Factorization 将多个小型矩阵乘法合并为一个大型融合操作。 | ✅ | ➖ | ➖ |
| `enhancer`   | 通过应用去噪或放大等后处理算法来增强模型输出。 | ❌ | - | ✅ |

✅（改善），➖（大致相同），❌（恶化）

在 [Pruna 文档](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms)中了解完整的优化方法。

## 安装

使用以下命令安装 Pruna。

```bash
pip install pruna
```

## 优化 Diffusers 模型

如下所示，Diffusers 模型支持多种优化算法。

    
以下示例通过组合使用 factorizer、compiler 和 cacher 算法来优化 [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)。该组合可将推理速度提升最高 4.2 倍，并将 GPU 峰值内存占用从 34.7GB 降低至 28.0GB，同时几乎保持相同的输出质量。

> [!TIP]
> 请参阅 [Pruna 优化](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html) 文档，以了解本示例中使用的优化技术。

    
首先，定义一个 `SmashConfig`，指定要使用的优化算法。要优化模型，请使用 `smash` 将 pipeline 和 `SmashConfig` 包装起来，然后像往常一样使用 pipeline 进行推理。

```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]
```

    
优化完成后，我们可以使用 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")
```

## 评估和基准测试 Diffusers 模型

Pruna 提供了 [EvaluationAgent](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)，用于评估您优化后模型的质量。

我们可以关注诸如总耗时和吞吐量等指标，以及用于评估的数据集。我们可以定义一个模型并将其传递给 `EvaluationAgent`。

我们可以使用 `EvaluationAgent` 加载并评估优化后的模型，并将其传递给 `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")
```

除了将优化后的模型与基础模型进行比较，您还可以单独评估 `diffusers` 模型。如果您希望在未应用优化的情况下评估模型性能，这将非常有用。我们可以使用 `PrunaModel` 包装器，并在其上运行 `EvaluationAgent` 来实现这一点。

```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)
```

现在你已经了解了如何优化和评估模型，可以开始使用 Pruna 来优化自己的模型了。幸运的是，我们提供了许多示例来帮助你入门。

> [!TIP]
> 有关 Flux 基准测试的更多详细信息，请参阅 [Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)!](https://huggingface.co/blog/PrunaAI/flux-fastest-image-generation-endpoint) 博客文章以及 [InferBench](https://huggingface.co/spaces/PrunaAI/InferBench) Space。

## 参考

- [Pruna](https://github.com/PrunaAI/pruna)
- [Pruna 优化](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms)
- [Pruna 评估](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)
- [Pruna 教程](https://docs.pruna.ai/en/stable/docs_pruna/tutorials/index.html)

Xet Storage Details

Size:: 7.31 kB
Xet hash:: 61c0997f8b5b01ff063cc80fef911098931ba82b2ca38cf356fbe3f714482fb9

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# Pruna

	[Pruna](https://github.com/PrunaAI/pruna) 是一个模型优化框架，提供多种优化方法——量化、剪枝、缓存、编译——以加速推理并降低内存占用。下图展示了这些优化方法的总体概览。

	\| 技术 \| 描述 \| 速度 \| 内存 \| 质量 \|
	\|--------------\|-----------------------------------------------------------------------------------------------\|:-----:\|:------:\|:-------:\|
	\| `batcher` \| 将多个输入组合在一起以同时处理，从而提高计算效率并缩短处理时间。 \| ✅ \| ❌ \| ➖ \|
	\| `cacher` \| 存储计算的中间结果，以加速后续操作。 \| ✅ \| ➖ \| ➖ \|
	\| `compiler` \| 针对特定硬件优化模型。 \| ✅ \| ➖ \| ➖ \|
	\| `distiller` \| 训练一个更小、更简单的模型，以模拟一个更大、更复杂的模型。 \| ✅ \| ✅ \| ❌ \|
	\| `quantizer` \| 降低权重和激活值的精度，从而减少内存需求。 \| ✅ \| ✅ \| ❌ \|
	\| `pruner` \| 移除不太重要或冗余的连接和神经元，从而得到一个更稀疏、更高效的网络。 \| ✅ \| ✅ \| ❌ \|
	\| `recoverer` \| 恢复模型压缩后的性能。 \| ➖ \| ➖ \| ✅ \|
	\| `factorizer` \| Factorization 将多个小型矩阵乘法合并为一个大型融合操作。 \| ✅ \| ➖ \| ➖ \|
	\| `enhancer` \| 通过应用去噪或放大等后处理算法来增强模型输出。 \| ❌ \| - \| ✅ \|

	✅（改善），➖（大致相同），❌（恶化）

	在 [Pruna 文档](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms)中了解完整的优化方法。

	## 安装

	使用以下命令安装 Pruna。

	```bash
	pip install pruna
	```

	## 优化 Diffusers 模型

	如下所示，Diffusers 模型支持多种优化算法。



	以下示例通过组合使用 factorizer、compiler 和 cacher 算法来优化 [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)。该组合可将推理速度提升最高 4.2 倍，并将 GPU 峰值内存占用从 34.7GB 降低至 28.0GB，同时几乎保持相同的输出质量。

	> [!TIP]
	> 请参阅 [Pruna 优化](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html) 文档，以了解本示例中使用的优化技术。



	首先，定义一个 `SmashConfig`，指定要使用的优化算法。要优化模型，请使用 `smash` 将 pipeline 和 `SmashConfig` 包装起来，然后像往常一样使用 pipeline 进行推理。

	```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]
	```



	优化完成后，我们可以使用 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")
	```

	## 评估和基准测试 Diffusers 模型

	Pruna 提供了 [EvaluationAgent](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)，用于评估您优化后模型的质量。

	我们可以关注诸如总耗时和吞吐量等指标，以及用于评估的数据集。我们可以定义一个模型并将其传递给 `EvaluationAgent`。

	我们可以使用 `EvaluationAgent` 加载并评估优化后的模型，并将其传递给 `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")
	```

	除了将优化后的模型与基础模型进行比较，您还可以单独评估 `diffusers` 模型。如果您希望在未应用优化的情况下评估模型性能，这将非常有用。我们可以使用 `PrunaModel` 包装器，并在其上运行 `EvaluationAgent` 来实现这一点。

	```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)
	```

	现在你已经了解了如何优化和评估模型，可以开始使用 Pruna 来优化自己的模型了。幸运的是，我们提供了许多示例来帮助你入门。

	> [!TIP]
	> 有关 Flux 基准测试的更多详细信息，请参阅 [Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)!](https://huggingface.co/blog/PrunaAI/flux-fastest-image-generation-endpoint) 博客文章以及 [InferBench](https://huggingface.co/spaces/PrunaAI/InferBench) Space。

	## 参考

	- [Pruna](https://github.com/PrunaAI/pruna)
	- [Pruna 优化](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms)
	- [Pruna 评估](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)
	- [Pruna 教程](https://docs.pruna.ai/en/stable/docs_pruna/tutorials/index.html)