NVFP4

Lượng tử hóa NVFP4 đóng gói các trọng số tuyến tính độ chính xác đầy đủ vào định dạng số thực 4-bit của NVIDIA trong khi mô hình được nạp. NVFP4Config thay thế các module torch.nn.Linear đủ điều kiện, không có bias, có cả in_features và out_features đều chia hết cho 16, bằng một triển khai tuyến tính NVFP4 từ NVFP4 Hub kernel. Các giao diện attention và MLP của mô hình không được thay thế.

NVFP4 yêu cầu GPU Blackwell có khả năng tính toán 10.0 hoặc mới hơn, bản build PyTorch hỗ trợ CUDA tương thích và gói kernels.

Cài đặt Accelerate và phiên bản tương thích của kernels.

pip install --upgrade accelerate kernels

Truyền NVFP4Config vào from_pretrained() với một thiết bị CUDA duy nhất. Trọng số được lượng tử hóa khi chúng được nạp, vì vậy checkpoint nguồn phải chứa các trọng số dạng số dấu phẩy động.

import torch

from transformers import AutoModelForCausalLM, AutoTokenizer, NVFP4Config


model_id = "meta-llama/Llama-3.2-1B"
quantization_config = NVFP4Config()
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    device_map="cuda",
    quantization_config=quantization_config,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("NVFP4 is", return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Sử dụng modules_to_not_convert để giữ các module được chọn ở độ chính xác ban đầu.

quantization_config = NVFP4Config(modules_to_not_convert=["vision", "lm_head"])

Các module linear NVFP4 hỗ trợ torch.compile. Lệnh gọi đầu tiên sau khi biên dịch bao gồm thời gian biên dịch đồ thị, vì vậy hãy khởi động nóng mô hình trước khi đo thông lượng sinh.

Hạn chế hiện tại

• Chỉ hỗ trợ một thiết bị CUDA. Các cấu hình song song tensor và device_map đa thiết bị sẽ bị từ chối cho đến khi hành vi chia nhỏ (sharding) của metadata tỷ lệ NVFP4 được xác định.
• CPU và disk offload không được hỗ trợ.
• Các checkpoint NVFP4 đã lượng tử hóa trước không được hỗ trợ.
• Các mô hình NVFP4 hiện tại không thể được serialize bằng save_pretrained() hoặc được huấn luyện.