NVIDIA ModelOpt

NVIDIA-ModelOpt est une bibliothèque unifiée de techniques d’optimisation de modèles de pointe, telles que la quantification, l’élagage, la distillation, le décodage spéculatif, etc. Elle compresse les modèles d’apprentissage profond pour les frameworks de déploiement en aval tels que TensorRT-LLM ou TensorRT afin d’optimiser la vitesse d’inférence.

Avant de commencer, assurez-vous que nvidia_modelopt est installé.

pip install -U "nvidia_modelopt[hf]"

Quantifier un modèle en passant NVIDIAModelOptConfig à from_pretrained() (vous pouvez également charger des modèles pré-quantifiés). Cela fonctionne pour tout modèle de toute modalité, à condition qu’il prenne en charge le chargement avec Accelerate et qu’il contienne des couches torch.nn.Linear.

L'exemple ci-dessous ne quantise que les poids en 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")

Remarque :

Les méthodes de quantification de NVIDIA-ModelOpt sont conçues pour réduire l’empreinte mémoire des poids du modèle à l’aide de diverses techniques de QAT (Quantization-Aware Training) et de PTQ (Post-Training Quantization), tout en préservant les performances du modèle. Toutefois, le gain de performance réel lors de l’inférence dépend du framework de déploiement (par exemple, TRT-LLM, TensorRT) et de la configuration matérielle spécifique.

Plus de détails peuvent être trouvés ici.

NVIDIAModelOptConfig

La classe NVIDIAModelOptConfig accepte trois paramètres :

• quant_type : une valeur de type chaîne mentionnant l'un des types de quantification ci-dessous.
• modules_to_not_convert : une liste de noms complets ou partiels de modules pour lesquels la quantification ne doit pas être effectuée. Par exemple, pour ne quantifier aucun des blocs de projection pos_embed du SD3Transformer2DModel, on spécifierait : modules_to_not_convert=["pos_embed.proj.weight"].
• disable_conv_quantization : une valeur booléenne qui, lorsqu’elle est définie sur True, désactive la quantification pour toutes les couches convolutives du modèle. Cela est utile car la quantification par canal et par bloc ne fonctionne généralement pas bien avec les couches convolutives (utilisée avec INT4, NF4, NVFP4). Si vous souhaitez désactiver la quantification pour des couches convolutives spécifiques, utilisez plutôt modules_to_not_convert.
• algorithm : l’algorithme à utiliser pour déterminer l’échelle, par défaut "max". Consultez la documentation de modelopt pour plus d’algorithmes et de détails.
• forward_loop : fonction de la boucle avant à utiliser pour calibrer l’activation lors de la quantification. Si elle n’est pas fournie, elle s’appuie sur des valeurs d’échelle statiques calculées à partir des poids uniquement.
• kwargs : un dictionnaire d’arguments nommés à passer à la méthode de quantification sous-jacente, qui sera appelée en fonction de quant_type.

Types de quantification pris en charge

ModelOpt prend en charge la quantification int8, fp8, int4, nf4 et nvfp4, au niveau des poids uniquement, par canal et par bloc. Les méthodes de quantification sont conçues pour réduire l’empreinte mémoire des poids du modèle tout en maintenant ses performances lors de l’inférence.

La quantification des poids uniquement stocke les poids du modèle dans un type de données à faible nombre de bits, mais effectue les calculs avec un type de données de plus haute précision, comme bfloat16. Cela réduit les besoins en mémoire liés aux poids du modèle, mais conserve les pics de mémoire pour le calcul des activations.

Les méthodes de quantisation prises en charge sont les suivantes :

Reportez-vous à la documentation officielle de modelopt pour une meilleure compréhension des méthodes de quantification disponibles et de la liste exhaustive des options de configuration.

Sérialisation et désérialisation des modèles quantifiés

Pour sérialiser un modèle quantifié dans un dtype donné, chargez d'abord le modèle avec le dtype de quantification souhaité, puis enregistrez-le à l'aide de la méthode 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)

Pour charger un modèle quantifié sérialisé, utilisez la méthode 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")