NVIDIA ModelOpt

NVIDIA-ModelOpt é uma biblioteca unificada de técnicas de otimização de modelos de última geração, como quantização, poda, destilação, decodificação especulativa, etc. Ela comprime modelos de aprendizado profundo para frameworks de implantação downstream, como TensorRT-LLM ou TensorRT, a fim de otimizar a velocidade de inferência.

Antes de começar, verifique se você tem o nvidia_modelopt instalado.

pip install -U "nvidia_modelopt[hf]"

Quantize um modelo passando NVIDIAModelOptConfig para from_pretrained() (você também pode carregar modelos pré-quantizados). Isso funciona para qualquer modelo em qualquer modalidade, desde que ele suporte carregamento com Accelerate e contenha camadas torch.nn.Linear.

O exemplo abaixo apenas quantiza os pesos para 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")

Observação:

Os métodos de quantização no NVIDIA-ModelOpt são projetados para reduzir a pegada de memória dos pesos do modelo usando diversas técnicas de QAT (Quantization-Aware Training) e PTQ (Post-Training Quantization), mantendo o desempenho do modelo. No entanto, o ganho de desempenho real durante a inferência depende do framework de implantação (por exemplo, TRT-LLM, TensorRT) e da configuração de hardware específica.

Mais detalhes podem ser encontrados aqui.

NVIDIAModelOptConfig

A classe NVIDIAModelOptConfig aceita três parâmetros:

• quant_type: Um valor de string que menciona um dos tipos de quantização abaixo.
• modules_to_not_convert: Uma lista de nomes completos ou parciais de módulos para os quais a quantização não deve ser realizada. Por exemplo, para não quantizar os blocos de projeção pos_embed do SD3Transformer2DModel, deve-se especificar: modules_to_not_convert=["pos_embed.proj.weight"].
• disable_conv_quantization: Um valor booleano que, quando definido como True, desativa a quantização para todas as camadas convolucionais no modelo. Isso é útil, pois a quantização por canal e por bloco geralmente não funciona bem com camadas convolucionais (usado com INT4, NF4, NVFP4). Se você deseja desativar a quantização para camadas convolucionais específicas, use modules_to_not_convert em vez disso.
• algorithm: O algoritmo a ser usado para determinar a escala, padrão é "max". Você pode consultar a documentação do modelopt para mais algoritmos e detalhes.
• forward_loop: A função do loop de forward a ser usada para calibrar a ativação durante a quantização. Se não for fornecida, ela se baseia em valores de escala estáticos calculados usando apenas os pesos.
• kwargs: Um dicionário de argumentos por palavra-chave para passar ao método de quantização subjacente, que será chamado com base em quant_type.

Tipos de quantização suportados

O ModelOpt suporta quantização int8, fp8, int4, nf4 e nvfp4 apenas de pesos, por canal e por bloco. Os métodos de quantização são projetados para reduzir a pegada de memória dos pesos do modelo, mantendo o desempenho do modelo durante a inferência.

A quantização apenas de pesos armazena os pesos do modelo em um tipo de dados de baixa bit, mas executa o cálculo com um tipo de dados de maior precisão, como bfloat16. Isso reduz os requisitos de memória dos pesos do modelo, mas mantém os picos de memória para o cálculo de ativações.

Os métodos de quantização suportados são os seguintes:

Consulte a documentação oficial do modelopt para uma melhor compreensão dos métodos de quantização disponíveis e da lista exaustiva de opções de configuração disponíveis.

Serialização e desserialização de modelos quantizados

Para serializar um modelo quantizado em um dtype específico, primeiro carregue o modelo com o dtype de quantização desejado e, em seguida, salve-o usando o método 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)

Para carregar um modelo quantizado serializado, use o método 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")