bitsandbytes

bitsandbytes — самый простой вариант для квантования модели до 8 и 4 бит. 8-битное квантование умножает выбросы в fp16 на значения, не являющиеся выбросами, в int8, затем преобразует значения, не являющиеся выбросами, обратно в fp16 и складывает их, чтобы вернуть веса в fp16. Это снижает негативное влияние выбросов на производительность модели.

4-битовая квантизация сжимает модель ещё сильнее и часто используется вместе с QLoRA для дообучения квантованных LLM.

В этом руководстве показано, как квантование позволяет запускать FLUX.1-dev на видеокарте с менее чем 16 ГБ видеопамяти, а также на бесплатном экземпляре Google Colab.

Для использования bitsandbytes убедитесь, что у вас установлены следующие библиотеки:

pip install diffusers transformers accelerate bitsandbytes -U

Теперь вы можете квантовать модель, передав BitsAndBytesConfig в from_pretrained(). Это работает для любой модели в любой модальности, если она поддерживает загрузку с помощью Accelerate и содержит слои torch.nn.Linear.

Квантование модели в 8 бит сокращает потребление памяти вдвое:

bitsandbytes поддерживается как в Transformers, так и в Diffusers, поэтому вы можете квантовать и FluxTransformer2DModel, и T5EncoderModel.

Для GPU серии Ada и более новых мы рекомендуем изменить dtype на torch.bfloat16.

Модели CLIPTextModel и AutoencoderKL не квантуются, так как они и так имеют небольшой размер, а AutoencoderKL содержит лишь несколько слоёв torch.nn.Linear.

Скопировано

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)

text_encoder_2_8bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

По умолчанию все остальные модули, такие как torch.nn.LayerNorm, преобразуются в torch.float16. Вы можете изменить тип данных этих модулей с помощью параметра dtype.

Скопировано

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
+   dtype=torch.float32,
)

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

Установка device_map="auto" автоматически заполняет сначала всё доступное пространство на GPU, затем на CPU и, наконец, на жёстком диске (самый медленный вариант), если памяти всё ещё недостаточно.

Скопировано

from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer_8bit,
    text_encoder_2=text_encoder_2_8bit,
    dtype=torch.float16,
    device_map="auto",
)

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
    "max_sequence_length": 512,
}

image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]

Если памяти достаточно, вы также можете напрямую переместить конвейер на GPU с помощью .to("cuda") и применить enable_model_cpu_offload() для оптимизации использования памяти GPU.

Once a model is quantized, you can push the model to the Hub with the push_to_hub() method. The quantization config.json file is pushed first, followed by the quantized model weights. You can also save the serialized 8-bit models locally with save_pretrained().

Обучение с весами 8-бит и 4-бит поддерживается только для обучения дополнительных параметров.

Проверьте потребление памяти с помощью метода get_memory_footprint:

print(model.get_memory_footprint())

Обратите внимание, что это показывает только объём памяти, занимаемый параметрами модели, и не оценивает требования к памяти при инференсе.

Квантованные модели можно загрузить с помощью метода from_pretrained() без необходимости указывать параметры quantization_config:

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

model_4bit = AutoModel.from_pretrained(
    "hf-internal-testing/flux.1-dev-nf4-pkg", subfolder="transformer"
)

8-бит (алгоритм LLM.int8())

Подробнее о деталях 8-битного квантования читайте в этой статье в блоге!

В этом разделе рассматриваются некоторые специфические особенности 8-битных моделей, такие как пороги выбросов и пропуск конвертации модулей.

Порог выбросов

«Аутлайер» — это значение скрытого состояния, превышающее определённый порог, и такие значения вычисляются в fp16. Как правило, значения распределены нормально ([-3.5, 3.5]), однако для больших моделей это распределение может сильно отличаться ([-60, 6] или [6, 60]). 8-битное квантование хорошо работает для значений ~5, но при превышении этого диапазона наблюдается значительное снижение производительности. Хорошим значением порога по умолчанию является 6, однако для более нестабильных моделей (маленькие модели или дообучение) может потребоваться более низкий порог.

Чтобы подобрать оптимальный порог для вашей модели, мы рекомендуем экспериментировать с параметром llm_int8_threshold в BitsAndBytesConfig:

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True, llm_int8_threshold=10,
)

model_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quantization_config,
)

Пропустить конвертацию модуля

Для некоторых моделей не обязательно квантовать каждый модуль до 8 бит, так как это может привести к нестабильности. Например, для диффузионных моделей, таких как Stable Diffusion 3, модуль proj_out можно пропустить с помощью параметра llm_int8_skip_modules в BitsAndBytesConfig:

from diffusers import SD3Transformer2DModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True, llm_int8_skip_modules=["proj_out"],
)

model_8bit = SD3Transformer2DModel.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    subfolder="transformer",
    quantization_config=quantization_config,
)

4-бит (алгоритм QLoRA)

Подробнее о деталях можно узнать в этой статье в блоге.

В этом разделе рассматриваются некоторые специфические особенности 4-битных моделей, такие как изменение типа данных для вычислений, использование типа данных Normal Float 4 (NF4) и применение вложенной квантизации.

Тип вычислительных данных

Чтобы ускорить вычисления, вы можете изменить тип данных с float32 (значение по умолчанию) на bf16, используя параметр bnb_4bit_compute_dtype в BitsAndBytesConfig:

import torch
from diffusers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

Обычный Float 4 (NF4)

NF4 — это 4-битный тип данных из статьи QLoRA, адаптированный для весов, инициализированных из нормального распределения. NF4 следует использовать для обучения базовых моделей с 4-битным квантованием. Настройка выполняется с помощью параметра bnb_4bit_quant_type в BitsAndBytesConfig:

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Для инференса параметр bnb_4bit_quant_type не оказывает значительного влияния на производительность. Однако, чтобы сохранить согласованность с весами модели, следует использовать значения bnb_4bit_compute_dtype и dtype.

Вложенная квантизация

Вложенная квантизация — это техника, позволяющая сэкономить дополнительную память без потери производительности. Эта функция выполняет вторую квантизацию уже квантованных весов, что позволяет сэкономить дополнительные 0,4 бита на параметр.

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Декуантизация моделей bitsandbytes

После квантования вы можете деквантовать модель до её исходной точности, однако это может привести к небольшой потере качества. Убедитесь, что у вас достаточно GPU RAM для размещения деквантованной модели.

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

text_encoder_2_4bit.dequantize()
transformer_4bit.dequantize()

torch.compile

Ускорьте инференс с помощью torch.compile. Убедитесь, что у вас установлена последняя версия bitsandbytes, и мы также рекомендуем установить PyTorch nightly.

torch._dynamo.config.capture_dynamic_output_shape_ops = True

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)

quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True)
transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)

На RTX 4090 с компиляцией генерация 4-битного Flux заняла 25,809 секунды против 32,570 секунд без неё.

Подробнее о том, как это работает, можно узнать в скрипте бенчмаркинга.

Ресурсы

• Ноутбук с примером вывода модели Flux.1 Dev в бесплатном Colab
• Обучение