Начало работы

Квантование направлено на представление данных с меньшим количеством битов при сохранении точности исходных данных. Это часто означает преобразование типа данных для представления той же информации с меньшим числом битов. Например, если веса вашей модели хранятся в формате 32-битных чисел с плавающей запятой и квантуются до 16-битных чисел с плавающей запятой, размер модели уменьшается вдвое, что облегчает её хранение и снижает потребление памяти. Более низкая точность также может ускорить инференс, поскольку выполнение вычислений с меньшим числом битов занимает меньше времени.

Diffusers поддерживает несколько бэкендов квантования, чтобы сделать крупные диффузионные модели, такие как Flux, более доступными. В этом руководстве показано, как использовать класс PipelineQuantizationConfig для квантования пайплайна при его инициализации из предобученной или неквантованной контрольной точки.

Квантование на уровне пайплайна

Существует два способа использования PipelineQuantizationConfig в зависимости от того, насколько глубоко вы хотите настроить конфигурацию квантования.

• для базовых сценариев использования определите аргументы quant_backend, quant_kwargs и components_to_quantize
• для детального управления квантованием определите quant_mapping, который предоставляет конфигурацию квантования для отдельных компонентов модели

Базовая квантизация

Инициализируйте PipelineQuantizationConfig со следующими параметрами.

• quant_backend определяет, какой бэкенд квантования использовать. В настоящее время поддерживаются следующие бэкенды: bitsandbytes_4bit, bitsandbytes_8bit, gguf, quanto и torchao.
• quant_kwargs задаёт аргументы квантования, которые следует использовать.

Аргументы quant_kwargs различаются в зависимости от бэкенда. Чтобы просмотреть аргументы для каждого бэкенда, обратитесь к документации Quantization API.

• components_to_quantize определяет, какие компоненты пайплайна следует квантовать. Как правило, квантовать нужно наиболее вычислительно затратные компоненты, например трансформер. Если в пайплайне есть несколько текстовых энкодеров, как в FluxPipeline, стоит рассмотреть квантование и их. В примере ниже квантуется текстовый энкодер T5 в FluxPipeline, при этом модель CLIP остаётся без изменений.

  components_to_quantize принимает список для нескольких моделей или строку для одной модели.

Пример ниже загружает бэкенд bitsandbytes со следующими аргументами из BitsAndBytesConfig: load_in_4bit, bnb_4bit_quant_type и bnb_4bit_compute_dtype.

import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_backend="bitsandbytes_4bit",
    quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
    components_to_quantize=["transformer", "text_encoder_2"],
)

Передайте pipeline_quant_config в from_pretrained(), чтобы квантовать пайплайн.

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

image = pipe("photo of a cute dog").images[0]

Расширенная квантизация

Аргумент quant_mapping предоставляет дополнительные варианты для квантования каждого отдельного компонента в конвейере, например, для объединения различных бэкендов квантования.

Инициализируйте PipelineQuantizationConfig и передайте в него quant_mapping. С помощью quant_mapping вы можете указать параметры квантования для каждого компонента конвейера, например, для трансформера и текстового энкодера.

Пример ниже использует два бэкенда квантования, QuantoConfig и transformers.BitsAndBytesConfig, для трансформера и текстового энкодера.

import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers.quantization_config import QuantoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={
        "transformer": QuantoConfig(weights_dtype="int8"),
        "text_encoder_2": TransformersBitsAndBytesConfig(
            load_in_4bit=True, compute_dtype=torch.bfloat16
        ),
    }
)

В Transformers есть отдельный backend bitsandbytes. Для компонентов, поступающих из Transformers, необходимо импортировать и использовать transformers.BitsAndBytesConfig. Например, text_encoder_2 в FluxPipeline является T5EncoderModel из Transformers, поэтому нужно использовать transformers.BitsAndBytesConfig вместо diffusers.BitsAndBytesConfig.

Используйте метод basic quantization, описанный выше, если вы не хотите управлять этими отдельными импортами или не уверены, откуда берётся каждый компонент конвейера.

import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={
        "transformer": DiffusersBitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16),
        "text_encoder_2": TransformersBitsAndBytesConfig(
            load_in_4bit=True, compute_dtype=torch.bfloat16
        ),
    }
)

Передайте pipeline_quant_config в from_pretrained(), чтобы квантовать пайплайн.

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

image = pipe("photo of a cute dog").images[0]

Загрузка предквантованных контрольных точек

Некоторые бэкенды квантования поддерживают загрузку контрольных точек, которые уже были квантованы и загружены в Hub. В этом случае конфигурация квантования хранится в файле config.json модели и читается автоматически — вам не нужно создавать PipelineQuantizationConfig или передавать аргумент quantization_config.

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

Следующие бэкенды поддерживают загрузку предквантованных контрольных точек из коробки.

Ресурсы

Ознакомьтесь с материалами ниже, чтобы узнать больше о квантовании.

• Если вы только начинаете знакомство с квантованием, мы рекомендуем пройти следующие вводные курсы, созданные совместно с DeepLearning.AI.

  • Основы квантования с Hugging Face
  • Квантование: подробный разбор

• Если вы хотите добавить новый метод квантования, обратитесь к руководству по добавлению нового метода квантования.

• Обзор квантования в Transformers Overview содержит краткое описание преимуществ и недостатков различных бэкендов квантования.

• Ознакомьтесь с блогом Exploring Quantization Backends in Diffusers, где кратко описаны каждый из бэкендов квантования, принципы выбора бэкенда и способы сочетания квантования с другими методами оптимизации памяти.