NVIDIA ModelOpt

NVIDIA-ModelOpt — это единая библиотека передовых методов оптимизации моделей, таких как квантование, обрезка, дистилляция, спекулятивное декодирование и другие. Она сжимает модели глубокого обучения для фреймворков развертывания, таких как TensorRT-LLM или TensorRT, с целью оптимизации скорости вывода.

Перед началом убедитесь, что у вас установлен nvidia_modelopt.

pip install -U "nvidia_modelopt[hf]"

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

Приведённый ниже пример квантует только веса до FP8.

import torch
from diffusers import AutoModel, SanaPipeline, NVIDIAModelOptConfig

model_id = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
dtype = torch.bfloat16

quantization_config = NVIDIAModelOptConfig(quant_type="FP8", quant_method="modelopt")
transformer = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=dtype,
)
pipe = SanaPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=dtype,
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB")

prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")

Примечание:

Методы квантования в NVIDIA-ModelOpt разработаны для снижения объёма памяти, занимаемой весами модели, с использованием различных техник QAT (Quantization-Aware Training) и PTQ (Post-Training Quantization) при сохранении производительности модели. Однако фактический прирост производительности на этапе инференса зависит от фреймворка развёртывания (например, TRT-LLM, TensorRT) и конкретной аппаратной конфигурации.

Подробнее можно узнать здесь.

NVIDIAModelOptConfig

Класс NVIDIAModelOptConfig принимает три параметра:

• quant_type: строковое значение, указывающее один из следующих типов квантования.
• modules_to_not_convert: список полных или частичных имён модулей, для которых не следует выполнять квантование. Например, чтобы не выполнять квантование блоков проекции pos_embed модели SD3Transformer2DModel, нужно указать: modules_to_not_convert=["pos_embed.proj.weight"].
• disable_conv_quantization: логическое значение, которое при установке в True отключает квантование для всех свёрточных слоёв в модели. Это полезно, так как квантование по каналам и блокам, как правило, плохо работает со свёрточными слоями (используется с INT4, NF4, NVFP4). Если вы хотите отключить квантование для конкретных свёрточных слоёв, используйте вместо этого modules_to_not_convert.
• algorithm: Алгоритм, используемый для определения масштаба, по умолчанию "max". Дополнительные сведения об алгоритмах и деталях см. в документации modelopt.
• forward_loop: функция прямого прохода, используемая для калибровки активаций при квантовании. Если не задана, используются статические значения масштаба, вычисленные только на основе весов.
• kwargs: словарь именованных аргументов, передаваемых в базовый метод квантования, который будет вызван в зависимости от quant_type.

Поддерживаемые типы квантования

ModelOpt поддерживает квантование весов, каналов и блоков в форматах int8, fp8, int4, nf4 и nvfp4. Методы квантования разработаны для уменьшения объёма памяти, занимаемого весами модели, при сохранении её производительности во время инференса.

Квантование только весов хранит веса модели в определённом формате с низким числом бит, но выполняет вычисления в формате с более высокой точностью, например bfloat16. Это снижает требования к памяти, связанные с весами модели, но сохраняет пиковые значения памяти для вычислений активаций.

Поддерживаются следующие методы квантования:

См. официальную документацию modelopt, чтобы лучше понять доступные методы квантования и полный список доступных параметров конфигурации.

Сериализация и десериализация квантованных моделей

Чтобы сериализовать квантованную модель в заданном типе данных, сначала загрузите модель с нужным типом данных для квантования, а затем сохраните её с помощью метода save_pretrained().

import torch
from diffusers import AutoModel, NVIDIAModelOptConfig
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()

model_id = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
quant_config_fp8 = {"quant_type": "FP8", "quant_method": "modelopt"}
quant_config_fp8 = NVIDIAModelOptConfig(**quant_config_fp8)
model = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quant_config_fp8,
    dtype=torch.bfloat16,
)
model.save_pretrained('path/to/sana_fp8', safe_serialization=False)

Чтобы загрузить сериализованную квантованную модель, используйте метод from_pretrained().

import torch
from diffusers import AutoModel, NVIDIAModelOptConfig, SanaPipeline
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()

quantization_config = NVIDIAModelOptConfig(quant_type="FP8", quant_method="modelopt")
transformer = AutoModel.from_pretrained(
    "path/to/sana_fp8",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
pipe = SanaPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_600M_1024px_diffusers",
    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=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")