Buckets:
| # torchao | |
| [torchao](https://github.com/pytorch/ao) предоставляет высокопроизводительные dtypes и оптимизации на основе квантования и разреженности для инференса и обучения моделей PyTorch. Библиотека поддерживает любые модели в любой модальности, при условии, что они поддерживают загрузку с помощью [Accelerate](https://hf.co/docs/accelerate/index) и содержат слои `torch.nn.Linear`. | |
| Убедитесь, что Pytorch 2.5+ и torchao установлены с помощью команды ниже. | |
| ```bash | |
| uv pip install -U torch torchao | |
| ``` | |
| Каждый dtype квантования доступен как отдельный экземпляр класса [AOBaseConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.core.config.AOBaseConfig.html). Это обеспечивает более гибкие параметры конфигурации за счёт раскрытия большего числа доступных аргументов. | |
| Передайте `AOBaseConfig` для типа квантования, например [Int4WeightOnlyConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), в [TorchAoConfig](/docs/diffusers/v0.43.0/ru/api/quantization#diffusers.TorchAoConfig) внутри [from_pretrained()](/docs/diffusers/v0.43.0/ru/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"` квантует каждый слой на GPU по мере его загрузки. Это быстро, но временно требует дополнительной памяти GPU для исходных и квантованных весов. Если модель уже занимает большую часть памяти вашего GPU, загрузка может завершиться ошибкой нехватки памяти. В этом случае уберите `device_map` и переместите конвейер на GPU после загрузки: | |
| ```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" | |
| ``` | |
| Без `device_map` Diffusers квантует слои на CPU. Это медленнее, но позволяет избежать временного пика потребления памяти GPU во время квантования. Чтобы еще больше снизить использование памяти GPU, используйте вместо этого [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/ru/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload). Вы также можете квантовать дополнительные компоненты, например текстовый энкодер. | |
| ## torch.compile | |
| torchao поддерживает [torch.compile](../optimization/fp16#torchcompile), что позволяет ускорить инференс одной строкой кода. | |
| ```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) | |
| ``` | |
| См. эту [таблицу](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450) с бенчмарками скорости инференса и потребления памяти для Flux и CogVideoX. Дополнительные бенчмарки на различных аппаратных платформах также доступны в репозитории torchao [repository](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks). | |
| > [!TIP] | |
| > Схемы пост-тренировочного квантования FP8 в torchao эффективны для GPU с вычислительной мощностью не ниже 8.9 (RTX-4090, Hopper и т. д.). FP8 часто обеспечивает наилучший баланс между скоростью, потреблением памяти и качеством при генерации изображений и видео. Если ваш GPU совместим, мы рекомендуем использовать FP8 в сочетании с torch.compile. | |
| ## Поддерживаемые типы квантования | |
| torchao поддерживает квантование только весов, а также квантование весов и динамических активаций для int8, float3-float8 и uint1-uint7. | |
| Квантование только весов хранит веса модели в определённом формате с низким числом бит, но выполняет вычисления в формате с более высокой точностью, например `bfloat16`. Это снижает требования к памяти, связанные с весами модели, но сохраняет пиковые значения памяти для вычислений активаций. | |
| Динамическое квантование активаций хранит веса модели в формате с низким числом бит, одновременно квантуя активации на лету для экономии дополнительной памяти. Это снижает требования к памяти, связанные с весами модели, а также уменьшает накладные расходы памяти на вычисления активаций. Однако в некоторых случаях это может привести к снижению качества, поэтому рекомендуется тщательно тестировать различные модели. | |
| См. [официальную документацию torchao](https://docs.pytorch.org/ao/stable/index.html), чтобы лучше понять доступные методы квантования. Полный список вариантов конфигурации доступен [здесь](https://docs.pytorch.org/ao/main/workflows/inference.html#inference-workflows). | |
| Некоторые популярные конфигурации квантования приведены ниже: | |
| | **Категория** | **Классы конфигурации** | | |
| |---|---| | |
| | **Целочисленная квантизация** | [`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) | | |
| | **Квантование с плавающей точкой 8 бит** | [`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) | | |
| | **Квантование беззнаковых целых чисел** | [`IntxWeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.IntxWeightOnlyConfig.html) | | |
| ## Сериализация и десериализация квантованных моделей | |
| Чтобы сериализовать квантованную модель в заданном типе данных, сначала загрузите модель с нужным типом данных для квантования, а затем сохраните её с помощью метода [save_pretrained()](/docs/diffusers/v0.43.0/ru/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) | |
| ``` | |
| Чтобы загрузить сериализованную квантованную модель, используйте метод [from_pretrained()](/docs/diffusers/v0.43.0/ru/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") | |
| ``` | |
| Если вы используете `torch<=2.6.0`, некоторые методы квантования, такие как `uint4` weight-only, не могут быть загружены напрямую и могут вызывать `UnpicklingError` при попытке загрузки моделей, но работают как ожидается при их сохранении. Чтобы обойти эту проблему, можно вручную загрузить state dict в модель. Однако имейте в виду, что для этого необходимо использовать `weights_only=False` в `torch.load`, поэтому выполнять это следует только в том случае, если веса получены из надежного источника. | |
| ```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] | |
| > API [AutoModel](/docs/diffusers/v0.43.0/ru/api/models/auto_model#diffusers.AutoModel) поддерживается для PyTorch >= 2.6, как показано в примерах ниже. | |
| ## Ресурсы | |
| - [API квантования TorchAO](https://docs.pytorch.org/ao/stable/index.html) | |
| - [Примеры Diffusers-TorchAO](https://github.com/sayakpaul/diffusers-torchao) | |
Xet Storage Details
- Size:
- 12.6 kB
- Xet hash:
- 35d883ecfe6cdefdf8f45d6d94429948cc0fed57740636d0af90314eb7265533
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.