# torchao

[torchao](https://github.com/pytorch/ao) ofrece dtypes de alto rendimiento y optimizaciones basadas en cuantización y esparsidad para la inferencia y el entrenamiento de modelos de PyTorch. Es compatible con cualquier modelo de cualquier modalidad, siempre que admita la carga con [Accelerate](https://hf.co/docs/accelerate/index) y contenga capas `torch.nn.Linear`.

Asegúrate de que Pytorch 2.5+ y torchao estén instalados con el comando siguiente.

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

Cada dtype de cuantización está disponible como una instancia separada de la clase [AOBaseConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.core.config.AOBaseConfig.html). Esto ofrece opciones de configuración más flexibles al exponer más argumentos disponibles.

Pase el `AOBaseConfig` de un dtype de cuantización, como [Int4WeightOnlyConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), a [TorchAoConfig](/docs/diffusers/main/es/api/quantization#diffusers.TorchAoConfig) en [from_pretrained()](/docs/diffusers/main/es/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"` cuantiza cada capa en la GPU mientras se carga. Esto es rápido, pero requiere temporalmente memoria adicional de GPU para los pesos originales y cuantizados. Si el modelo ya utiliza la mayor parte de la memoria de tu GPU, la carga puede fallar con un error de memoria insuficiente. En ese caso, elimina `device_map` y mueve el pipeline a la GPU después de cargarlo:

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

Sin `device_map`, Diffusers cuantiza las capas en la CPU. Esto es más lento, pero evita el pico temporal de memoria de GPU durante la cuantización. Para reducir aún más el uso de memoria de GPU, usa [enable_model_cpu_offload()](/docs/diffusers/main/es/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) en su lugar. También puedes cuantizar componentes adicionales, como el codificador de texto.

## torch.compile

torchao admite [torch.compile](../optimization/fp16#torchcompile), que puede acelerar la inferencia con una sola línea 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 [tabla](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450) para los benchmarks de velocidad de inferencia y uso de memoria con Flux y CogVideoX. También hay más benchmarks en varios hardware disponibles en el [repositorio](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks) de torchao.

> [!TIP]
> Los esquemas de cuantización post-entrenamiento en FP8 de torchao son efectivos para GPUs con capacidad de cómputo de al menos 8.9 (RTX-4090, Hopper, etc.). FP8 suele ofrecer el mejor equilibrio entre velocidad, memoria y calidad al generar imágenes y videos. Recomendamos combinar FP8 y torch.compile si tu GPU es compatible.

## Tipos de cuantización admitidos

torchao admite la cuantización solo de pesos y la cuantización de pesos y activaciones dinámicas para int8, float3-float8 y uint1-uint7.

La cuantización solo de pesos almacena los pesos del modelo en un tipo de datos de bajo número de bits, pero realiza los cálculos con un tipo de datos de mayor precisión, como `bfloat16`. Esto reduce los requisitos de memoria derivados de los pesos del modelo, pero mantiene los picos de memoria para el cálculo de las activaciones.

La cuantización dinámica de activaciones almacena los pesos del modelo en un dtype de bajo número de bits, mientras que también cuantiza las activaciones sobre la marcha para ahorrar memoria adicional. Esto reduce los requisitos de memoria de los pesos del modelo, al tiempo que reduce la sobrecarga de memoria de los cálculos de activación. Sin embargo, esto puede implicar un compromiso en la calidad en algunos casos, por lo que se recomienda probar a fondo diferentes modelos.

Consulte la [documentación oficial de torchao](https://docs.pytorch.org/ao/stable/index.html) para comprender mejor los métodos de cuantización disponibles. Una lista exhaustiva de las opciones de configuración está disponible [aquí](https://docs.pytorch.org/ao/main/workflows/inference.html#inference-workflows).

Algunas configuraciones de cuantización populares y de ejemplo son las siguientes:

| **Categoría** | **Clases de configuración** |
|---|---|
| **Cuantización entera** | [`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) |
| **Cuantización de punto flotante 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) |
| **Cuantización de enteros sin signo** | [`IntxWeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.IntxWeightOnlyConfig.html) |

## Serialización y deserialización de modelos cuantizados

Para serializar un modelo cuantizado en un dtype dado, primero cargue el modelo con el dtype de cuantización deseado y luego guárdelo usando el método [save_pretrained()](/docs/diffusers/main/es/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 cargar un modelo cuantizado serializado, use el método [from_pretrained()](/docs/diffusers/main/es/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")
```

Si estás usando `torch<=2.6.0`, algunos métodos de cuantización, como `uint4` solo de pesos, no se pueden cargar directamente y pueden provocar un `UnpicklingError` al intentar cargar los modelos, aunque funcionan como se espera al guardarlos. Para evitar esto, se puede cargar el state dict manualmente en el modelo. Ten en cuenta, sin embargo, que esto requiere usar `weights_only=False` en `torch.load`, por lo que solo debe ejecutarse si los pesos se obtuvieron de una fuente de confianza.

```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]
> La API de [AutoModel](/docs/diffusers/main/es/api/models/auto_model#diffusers.AutoModel) es compatible con PyTorch >= 2.6, como se muestra en los ejemplos a continuación.

## Recursos

- [API de cuantización de TorchAO](https://docs.pytorch.org/ao/stable/index.html)
- [Ejemplos de Diffusers-TorchAO](https://github.com/sayakpaul/diffusers-torchao)