NVIDIA ModelOpt

NVIDIA-ModelOpt adalah pustaka terpadu yang berisi teknik optimasi model terkini seperti kuantisasi, pemangkasan, distilasi, dekode spekulatif, dan lain-lain. Pustaka ini mengompres model pembelajaran mendalam untuk kerangka kerja penerapan hilir seperti TensorRT-LLM atau TensorRT guna mengoptimalkan kecepatan inferensi.

Sebelum Anda memulai, pastikan Anda telah menginstal nvidia_modelopt.

pip install -U "nvidia_modelopt[hf]"

Kuantisasi model dengan meneruskan NVIDIAModelOptConfig ke from_pretrained() (Anda juga dapat memuat model yang sudah dikuantisasi). Ini berlaku untuk model apa pun dalam modalitas apa pun, selama model tersebut mendukung pemuatan dengan Accelerate dan berisi lapisan torch.nn.Linear.

Contoh di bawah ini hanya mengkuantisasi bobot ke 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")

Catatan:

Metode kuantisasi di NVIDIA-ModelOpt dirancang untuk mengurangi jejak memori bobot model menggunakan berbagai teknik QAT (Quantization-Aware Training) dan PTQ (Post-Training Quantization) sambil mempertahankan kinerja model. Namun, peningkatan kinerja aktual selama inferensi bergantung pada framework deployment (misalnya, TRT-LLM, TensorRT) dan konfigurasi hardware yang spesifik.

Detail lebih lanjut dapat ditemukan di sini.

NVIDIAModelOptConfig

Kelas NVIDIAModelOptConfig menerima tiga parameter:

• quant_type: Nilai string yang menyebutkan salah satu tipe kuantisasi di bawah ini.
• modules_to_not_convert: Daftar nama modul penuh/sebagian di mana kuantisasi tidak boleh dilakukan. Misalnya, untuk tidak melakukan kuantisasi pada blok proyeksi pos_embed dari SD3Transformer2DModel, Anda dapat menentukan: modules_to_not_convert=["pos_embed.proj.weight"].
• disable_conv_quantization: Nilai boolean yang jika diatur ke True akan menonaktifkan kuantisasi untuk semua lapisan konvolusional dalam model. Ini berguna karena kuantisasi channel dan block umumnya tidak bekerja dengan baik pada lapisan konvolusional (dipakai dengan INT4, NF4, NVFP4). Jika Anda ingin menonaktifkan kuantisasi untuk lapisan konvolusional tertentu, gunakan modules_to_not_convert sebagai gantinya.
• algorithm: Algoritma yang digunakan untuk menentukan skala, secara default "max". Anda dapat memeriksa dokumentasi modelopt untuk algoritma dan detail lainnya.
• forward_loop: Fungsi loop forward yang digunakan untuk mengkalibrasi aktivasi selama kuantisasi. Jika tidak disediakan, ia mengandalkan nilai skala statis yang dihitung menggunakan hanya bobot.
• kwargs: Dict berisi argumen kata kunci untuk diteruskan ke metode kuantisasi dasar yang akan dipanggil berdasarkan quant_type.

Jenis kuantisasi yang didukung

ModelOpt mendukung kuantisasi int8, fp8, int4, nf4, dan nvfp4 untuk weight-only, channel, dan block. Metode kuantisasi ini dirancang untuk mengurangi jejak memori dari bobot model sambil mempertahankan kinerja model selama inferensi.

Kuantisasi hanya-berat menyimpan bobot model dalam tipe data bit rendah tertentu tetapi melakukan komputasi dengan tipe data presisi yang lebih tinggi, seperti bfloat16. Hal ini menurunkan kebutuhan memori dari bobot model tetapi mempertahankan puncak memori untuk komputasi aktivasi.

Metode kuantisasi yang didukung adalah sebagai berikut:

Lihat dokumentasi modelopt resmi untuk pemahaman yang lebih baik tentang metode kuantisasi yang tersedia dan daftar lengkap opsi konfigurasi yang tersedia.

Menserialisasi dan Mendeserialisasi model terkuantisasi

Untuk mengserialisasi model terkuantisasi dalam dtype tertentu, pertama muat model dengan dtype kuantisasi yang diinginkan, lalu simpan menggunakan metode 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)

Untuk memuat model kuantisasi yang telah diserialisasi, gunakan metode 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")