NVIDIA ModelOpt

NVIDIA-ModelOpt là thư viện thống nhất các kỹ thuật tối ưu hóa mô hình tiên tiến như lượng tử hóa, cắt tỉa, chưng cất, giải mã suy đoán, v.v. Nó nén các mô hình học sâu cho các khung triển khai hạ nguồn như TensorRT-LLM hoặc TensorRT nhằm tối ưu hóa tốc độ suy luận.

Trước khi bắt đầu, hãy đảm bảo bạn đã cài đặt nvidia_modelopt.

pip install -U "nvidia_modelopt[hf]"

Thực hiện lượng tử hóa một mô hình bằng cách truyền NVIDIAModelOptConfig vào from_pretrained() (bạn cũng có thể tải các mô hình đã được lượng tử hóa trước). Cách này hoạt động với bất kỳ mô hình nào trong bất kỳ chế độ nào, miễn là nó hỗ trợ tải bằng Accelerate và chứa các lớp torch.nn.Linear.

Ví dụ bên dưới chỉ lượng tử hóa trọng số xuống 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")

Lưu ý:

Các phương pháp lượng tử hóa trong NVIDIA-ModelOpt được thiết kế để giảm dung lượng bộ nhớ của trọng số mô hình bằng cách sử dụng các kỹ thuật QAT (Quantization-Aware Training) và PTQ (Post-Training Quantization) khác nhau, đồng thời duy trì hiệu suất của mô hình. Tuy nhiên, mức cải thiện hiệu suất thực tế trong quá trình suy luận phụ thuộc vào khung triển khai (ví dụ: TRT-LLM, TensorRT) và cấu hình phần cứng cụ thể.

Chi tiết hơn có thể được tìm thấy tại đây.

NVIDIAModelOptConfig

Lớp NVIDIAModelOptConfig nhận ba tham số:

• quant_type: Một giá trị chuỗi đề cập đến một trong các kiểu lượng tử hóa bên dưới.
• modules_to_not_convert: Một danh sách các tên module đầy đủ/đầy đủ một phần mà không nên thực hiện lượng tử hóa. Ví dụ, để không thực hiện bất kỳ lượng tử hóa nào cho các khối chiếu pos_embed của SD3Transformer2DModel, bạn sẽ chỉ định: modules_to_not_convert=["pos_embed.proj.weight"].
• disable_conv_quantization: Một giá trị boolean, khi được đặt thành True sẽ tắt lượng tử hóa cho tất cả các lớp tích chập trong mô hình. Điều này hữu ích vì lượng tử hóa theo kênh và theo khối thường không hoạt động tốt với các lớp tích chập (được sử dụng với INT4, NF4, NVFP4). Nếu bạn muốn tắt lượng tử hóa cho các lớp tích chập cụ thể, hãy sử dụng modules_to_not_convert thay thế.
• algorithm: Thuật toán được sử dụng để xác định scale, mặc định là "max". Bạn có thể xem tài liệu modelopt để biết thêm về các thuật toán và chi tiết.
• forward_loop: Hàm vòng lặp forward được sử dụng để hiệu chỉnh kích hoạt trong quá trình lượng tử hóa. Nếu không được cung cấp, nó sẽ dựa vào các giá trị scale tĩnh được tính toán chỉ bằng trọng số.
• kwargs: Một dict chứa các đối số theo từ khóa để truyền cho phương pháp lượng tử hóa cơ bản, phương pháp này sẽ được gọi dựa trên quant_type.

Các kiểu lượng hóa được hỗ trợ

ModelOpt hỗ trợ lượng tử hóa int8, fp8, int4, nf4 và nvfp4 ở mức chỉ trọng số, kênh và khối. Các phương pháp lượng tử hóa được thiết kế để giảm dung lượng bộ nhớ của trọng số mô hình trong khi duy trì hiệu suất của mô hình trong quá trình suy luận.

Lượng tử hóa chỉ trọng số lưu trữ các trọng số của mô hình trong một kiểu dữ liệu bit thấp cụ thể nhưng thực hiện tính toán bằng một kiểu dữ liệu có độ chính xác cao hơn, chẳng hạn như bfloat16. Điều này làm giảm yêu cầu bộ nhớ từ các trọng số của mô hình nhưng vẫn giữ nguyên các đỉnh bộ nhớ cho việc tính toán kích hoạt.

Các phương pháp lượng tử hóa được hỗ trợ như sau:

Tham khảo tài liệu chính thức của modelopt để hiểu rõ hơn về các phương pháp lượng tử hóa có sẵn và danh sách đầy đủ các tùy chọn cấu hình.

Chuẩn hóa và khử chuẩn hóa các mô hình lượng tử hóa

Để serialize một model đã được lượng tử hóa trong một dtype cho trước, trước tiên hãy tải model với dtype lượng tử hóa mong muốn, sau đó lưu nó bằng phương thức 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)

Để tải một mô hình lượng tử hóa đã được serialize, hãy sử dụng phương thức 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")