Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /main /pt /quantization /torchao.md

mishig

8 days ago

preview code

download

raw

9.74 kB

# torchao

[torchao](https://github.com/pytorch/ao) fornece dtypes de alto desempenho e otimizações baseadas em quantização e esparsidade para inferência e treinamento de modelos PyTorch. É suportado para qualquer modelo em qualquer modalidade, desde que suporte carregamento com [Accelerate](https://hf.co/docs/accelerate/index) e contenha camadas `torch.nn.Linear`.

Certifique-se de que o Pytorch 2.5+ e o torchao estão instalados com o comando abaixo.

```bash
uv pip install -U torch torchao
```

Cada dtype de quantização está disponível como uma instância separada de uma classe [AOBaseConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.core.config.AOBaseConfig.html). Isso fornece opções de configuração mais flexíveis ao expor mais argumentos disponíveis.

Passe o `AOBaseConfig` de um dtype de quantização, como [Int4WeightOnlyConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), para [TorchAoConfig](/docs/diffusers/main/pt/api/quantization#diffusers.TorchAoConfig) em [from_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.from_pretrained).

```py
import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig
pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128, version=2))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)
```

`device_map="cuda"` quantiza cada camada na GPU durante o carregamento. Isso é rápido, mas requer temporariamente memória de GPU adicional para os pesos originais e quantizados. Se o modelo já usa a maior parte da memória da sua GPU, o carregamento pode falhar com um erro de falta de memória. Nesse caso, remova `device_map` e mova o pipeline para a GPU após o carregamento:

```py
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16,
)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"
```

Sem `device_map`, o Diffusers quantiza as camadas na CPU. Isso é mais lento, mas evita o pico temporário de memória de GPU durante a quantização. Para reduzir ainda mais o uso de memória de GPU, use [enable_model_cpu_offload()](/docs/diffusers/main/pt/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) no lugar. Você também pode quantizar componentes adicionais, como o text encoder.

## torch.compile

O torchao oferece suporte a [torch.compile](../optimization/fp16#torchcompile), que pode acelerar a inferência com uma única linha de código.

```python
import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int4WeightOnlyConfig
pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)
pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True)
```

Consulte esta [tabela](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450) para benchmarks de velocidade de inferência e uso de memória com Flux e CogVideoX. Mais benchmarks em diversos hardwares também estão disponíveis no [repositório](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks) do torchao.

> [!TIP]
> Os esquemas de quantização pós-treino em FP8 no torchao são eficazes para GPUs com capacidade de computação de pelo menos 8.9 (RTX-4090, Hopper, etc.). O FP8 frequentemente oferece o melhor equilíbrio entre velocidade, memória e qualidade ao gerar imagens e vídeos. Recomendamos combinar FP8 e torch.compile se sua GPU for compatível.

## Tipos de quantização suportados

O torchao suporta quantização apenas de pesos e quantização de pesos e ativações dinâmicas para int8, float3-float8 e uint1-uint7.

A quantização apenas de pesos armazena os pesos do modelo em um tipo de dados de baixa bit, mas executa o cálculo com um tipo de dados de maior precisão, como `bfloat16`. Isso reduz os requisitos de memória dos pesos do modelo, mas mantém os picos de memória para o cálculo de ativações.

A quantização dinâmica de ativações armazena os pesos do modelo em um dtype de baixo bit, quantizando também as ativações em tempo real para economizar memória adicional. Isso reduz os requisitos de memória dos pesos do modelo, além de reduzir a sobrecarga de memória das computações de ativação. No entanto, isso pode implicar em um compromisso de qualidade em alguns casos, portanto, é recomendado testar diferentes modelos minuciosamente.

Consulte a [documentação oficial do torchao](https://docs.pytorch.org/ao/stable/index.html) para uma melhor compreensão dos métodos de quantização disponíveis. Uma lista exaustiva das opções de configuração está disponível [aqui](https://docs.pytorch.org/ao/main/workflows/inference.html#inference-workflows).

Alguns exemplos de configurações populares de quantização são os seguintes:

| **Categoria** | **Classes de Configuração** |
|---|---|
| **Quantização de inteiros** | [`Int4WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), [`Int8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8WeightOnlyConfig.html), [`Int8DynamicActivationInt8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8DynamicActivationInt8WeightConfig.html) |
| **Quantização de ponto flutuante de 8 bits** | [`Float8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8WeightOnlyConfig.html), [`Float8DynamicActivationFloat8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8DynamicActivationFloat8WeightConfig.html) |
| **Quantização de inteiros sem sinal** | [`IntxWeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.IntxWeightOnlyConfig.html) |

## Serialização e desserialização de modelos quantizados

Para serializar um modelo quantizado em um dtype específico, primeiro carregue o modelo com o dtype de quantização desejado e, em seguida, salve-o usando o método [save_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.save_pretrained).

```python
import torch
from diffusers import AutoModel, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig
quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)
```

Para carregar um modelo quantizado serializado, use o método [from_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.from_pretrained).

```python
import torch
from diffusers import FluxPipeline, AutoModel
transformer = AutoModel.from_pretrained("/path/to/flux_int8wo", dtype=torch.bfloat16, use_safetensors=False)
pipe = FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
image.save("output.png")
```

Se você estiver usando `torch<=2.6.0`, alguns métodos de quantização, como `uint4` weight-only, não podem ser carregados diretamente e podem resultar em um `UnpicklingError` ao tentar carregar os modelos, mas funcionam como esperado ao salvá-los. Para contornar isso, é possível carregar o state dict manualmente no modelo. Observe, no entanto, que isso requer o uso de `weights_only=False` em `torch.load`, portanto, deve ser executado apenas se os pesos tiverem sido obtidos de uma fonte confiável.

```python
import torch
from accelerate import init_empty_weights
from diffusers import FluxPipeline, AutoModel, TorchAoConfig
from torchao.quantization import IntxWeightOnlyConfig
# Serialize the model
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=TorchAoConfig(IntxWeightOnlyConfig(dtype=torch.uint4)),
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_uint4wo", safe_serialization=False, max_shard_size="50GB")
# ...
# Load the model
state_dict = torch.load("/path/to/flux_uint4wo/diffusion_pytorch_model.bin", weights_only=False, map_location="cpu")
with init_empty_weights():
    transformer = AutoModel.from_config("/path/to/flux_uint4wo/config.json")
transformer.load_state_dict(state_dict, strict=True, assign=True)
```

> [!TIP]
> A API [AutoModel](/docs/diffusers/main/pt/api/models/auto_model#diffusers.AutoModel) é suportada para PyTorch >= 2.6, conforme mostrado nos exemplos abaixo.

## Recursos

- [API de Quantização do TorchAO](https://docs.pytorch.org/ao/stable/index.html)
- [Exemplos de Diffusers-TorchAO](https://github.com/sayakpaul/diffusers-torchao)

Xet Storage Details

Size:: 9.74 kB
Xet hash:: d27f152028c483f0d0dad3cef7c3da4e61cd99dae04d836e9e3453218694df06

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

	# torchao

	[torchao](https://github.com/pytorch/ao) fornece dtypes de alto desempenho e otimizações baseadas em quantização e esparsidade para inferência e treinamento de modelos PyTorch. É suportado para qualquer modelo em qualquer modalidade, desde que suporte carregamento com [Accelerate](https://hf.co/docs/accelerate/index) e contenha camadas `torch.nn.Linear`.

	Certifique-se de que o Pytorch 2.5+ e o torchao estão instalados com o comando abaixo.

	```bash
	uv pip install -U torch torchao
	```

	Cada dtype de quantização está disponível como uma instância separada de uma classe [AOBaseConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.core.config.AOBaseConfig.html). Isso fornece opções de configuração mais flexíveis ao expor mais argumentos disponíveis.

	Passe o `AOBaseConfig` de um dtype de quantização, como [Int4WeightOnlyConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), para [TorchAoConfig](/docs/diffusers/main/pt/api/quantization#diffusers.TorchAoConfig) em [from_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.from_pretrained).

	```py
	import torch
	from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
	from torchao.quantization import Int8WeightOnlyConfig

	pipeline_quant_config = PipelineQuantizationConfig(
	quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128, version=2))}
	)
	pipeline = DiffusionPipeline.from_pretrained(
	"black-forest-labs/FLUX.1-dev",
	quantization_config=pipeline_quant_config,
	dtype=torch.bfloat16,
	device_map="cuda" # or "mps", "xpu", "cpu"
	)
	```

	`device_map="cuda"` quantiza cada camada na GPU durante o carregamento. Isso é rápido, mas requer temporariamente memória de GPU adicional para os pesos originais e quantizados. Se o modelo já usa a maior parte da memória da sua GPU, o carregamento pode falhar com um erro de falta de memória. Nesse caso, remova `device_map` e mova o pipeline para a GPU após o carregamento:

	```py
	pipeline = DiffusionPipeline.from_pretrained(
	"black-forest-labs/FLUX.1-dev",
	quantization_config=pipeline_quant_config,
	torch_dtype=torch.bfloat16,
	)
	pipeline.to("cuda") # or "mps", "xpu", "cpu"
	```

	Sem `device_map`, o Diffusers quantiza as camadas na CPU. Isso é mais lento, mas evita o pico temporário de memória de GPU durante a quantização. Para reduzir ainda mais o uso de memória de GPU, use [enable_model_cpu_offload()](/docs/diffusers/main/pt/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) no lugar. Você também pode quantizar componentes adicionais, como o text encoder.

	## torch.compile

	O torchao oferece suporte a [torch.compile](../optimization/fp16#torchcompile), que pode acelerar a inferência com uma única linha de código.

	```python
	import torch
	from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
	from torchao.quantization import Int4WeightOnlyConfig

	pipeline_quant_config = PipelineQuantizationConfig(
	quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))}
	)
	pipeline = DiffusionPipeline.from_pretrained(
	"black-forest-labs/FLUX.1-dev",
	quantization_config=pipeline_quant_config,
	dtype=torch.bfloat16,
	device_map="cuda" # or "mps", "xpu", "cpu"
	)

	pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True)
	```

	Consulte esta [tabela](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450) para benchmarks de velocidade de inferência e uso de memória com Flux e CogVideoX. Mais benchmarks em diversos hardwares também estão disponíveis no [repositório](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks) do torchao.

	> [!TIP]
	> Os esquemas de quantização pós-treino em FP8 no torchao são eficazes para GPUs com capacidade de computação de pelo menos 8.9 (RTX-4090, Hopper, etc.). O FP8 frequentemente oferece o melhor equilíbrio entre velocidade, memória e qualidade ao gerar imagens e vídeos. Recomendamos combinar FP8 e torch.compile se sua GPU for compatível.

	## Tipos de quantização suportados

	O torchao suporta quantização apenas de pesos e quantização de pesos e ativações dinâmicas para int8, float3-float8 e uint1-uint7.

	A quantização apenas de pesos armazena os pesos do modelo em um tipo de dados de baixa bit, mas executa o cálculo com um tipo de dados de maior precisão, como `bfloat16`. Isso reduz os requisitos de memória dos pesos do modelo, mas mantém os picos de memória para o cálculo de ativações.

	A quantização dinâmica de ativações armazena os pesos do modelo em um dtype de baixo bit, quantizando também as ativações em tempo real para economizar memória adicional. Isso reduz os requisitos de memória dos pesos do modelo, além de reduzir a sobrecarga de memória das computações de ativação. No entanto, isso pode implicar em um compromisso de qualidade em alguns casos, portanto, é recomendado testar diferentes modelos minuciosamente.

	Consulte a [documentação oficial do torchao](https://docs.pytorch.org/ao/stable/index.html) para uma melhor compreensão dos métodos de quantização disponíveis. Uma lista exaustiva das opções de configuração está disponível [aqui](https://docs.pytorch.org/ao/main/workflows/inference.html#inference-workflows).

	Alguns exemplos de configurações populares de quantização são os seguintes:

	\| Categoria \| Classes de Configuração \|
	\|---\|---\|
	\| Quantização de inteiros \| [`Int4WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), [`Int8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8WeightOnlyConfig.html), [`Int8DynamicActivationInt8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8DynamicActivationInt8WeightConfig.html) \|
	\| Quantização de ponto flutuante de 8 bits \| [`Float8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8WeightOnlyConfig.html), [`Float8DynamicActivationFloat8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8DynamicActivationFloat8WeightConfig.html) \|
	\| Quantização de inteiros sem sinal \| [`IntxWeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.IntxWeightOnlyConfig.html) \|

	## Serialização e desserialização de modelos quantizados

	Para serializar um modelo quantizado em um dtype específico, primeiro carregue o modelo com o dtype de quantização desejado e, em seguida, salve-o usando o método [save_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.save_pretrained).

	```python
	import torch
	from diffusers import AutoModel, TorchAoConfig
	from torchao.quantization import Int8WeightOnlyConfig

	quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
	transformer = AutoModel.from_pretrained(
	"black-forest-labs/Flux.1-Dev",
	subfolder="transformer",
	quantization_config=quantization_config,
	dtype=torch.bfloat16,
	)
	transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)
	```

	Para carregar um modelo quantizado serializado, use o método [from_pretrained()](/docs/diffusers/main/pt/api/models/overview#diffusers.ModelMixin.from_pretrained).

	```python
	import torch
	from diffusers import FluxPipeline, AutoModel

	transformer = AutoModel.from_pretrained("/path/to/flux_int8wo", dtype=torch.bfloat16, use_safetensors=False)
	pipe = FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, dtype=torch.bfloat16)
	pipe.to("cuda") # or "mps", "xpu", "cpu"

	prompt = "A cat holding a sign that says hello world"
	image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
	image.save("output.png")
	```

	Se você estiver usando `torch<=2.6.0`, alguns métodos de quantização, como `uint4` weight-only, não podem ser carregados diretamente e podem resultar em um `UnpicklingError` ao tentar carregar os modelos, mas funcionam como esperado ao salvá-los. Para contornar isso, é possível carregar o state dict manualmente no modelo. Observe, no entanto, que isso requer o uso de `weights_only=False` em `torch.load`, portanto, deve ser executado apenas se os pesos tiverem sido obtidos de uma fonte confiável.

	```python
	import torch
	from accelerate import init_empty_weights
	from diffusers import FluxPipeline, AutoModel, TorchAoConfig
	from torchao.quantization import IntxWeightOnlyConfig

	# Serialize the model
	transformer = AutoModel.from_pretrained(
	"black-forest-labs/Flux.1-Dev",
	subfolder="transformer",
	quantization_config=TorchAoConfig(IntxWeightOnlyConfig(dtype=torch.uint4)),
	dtype=torch.bfloat16,
	)
	transformer.save_pretrained("/path/to/flux_uint4wo", safe_serialization=False, max_shard_size="50GB")
	# ...

	# Load the model
	state_dict = torch.load("/path/to/flux_uint4wo/diffusion_pytorch_model.bin", weights_only=False, map_location="cpu")
	with init_empty_weights():
	transformer = AutoModel.from_config("/path/to/flux_uint4wo/config.json")
	transformer.load_state_dict(state_dict, strict=True, assign=True)
	```

	> [!TIP]
	> A API [AutoModel](/docs/diffusers/main/pt/api/models/auto_model#diffusers.AutoModel) é suportada para PyTorch >= 2.6, conforme mostrado nos exemplos abaixo.

	## Recursos

	- [API de Quantização do TorchAO](https://docs.pytorch.org/ao/stable/index.html)
	- [Exemplos de Diffusers-TorchAO](https://github.com/sayakpaul/diffusers-torchao)