NVFP4

NVFP4 양자화는 모델이 로드되는 동안 전체 정밀도의 선형 가중치를 NVIDIA의 4비트 부동소수점 포맷으로 패키징합니다. NVFP4Config는 in_features와 out_features가 모두 16으로 나누어 떨어지는, 바이어 없는 torch.nn.Linear 모듈을 NVFP4 Hub 커널의 NVFP4 선형 구현으로 대체합니다. 모델의 어텐션 및 MLP 인터페이스는 대체되지 않습니다.

NVFP4는 compute capability 10.0 이상을 지원하는 Blackwell GPU, 호환되는 CUDA 지원 PyTorch 빌드, 그리고 kernels 패키지가 필요합니다.

Accelerate와 호환되는 버전의 kernels를 설치하세요.

pip install --upgrade accelerate kernels

단일 CUDA 디바이스를 사용하여 NVFP4Config를 from_pretrained()에 전달합니다. 가중치는 로드되는 시점에 양자화되므로, 소스 체크포인트에는 부동소수점 가중치가 포함되어 있어야 합니다.

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))

modules_to_not_convert를 사용하여 선택한 모듈을 원래 정밀도로 유지합니다.

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

NVFP4 linear 모듈은 torch.compile을 지원합니다. 첫 번째 컴파일된 호출에는 그래프 컴파일 시간이 포함되므로, 생성 처리량을 측정하기 전에 모델을 워밍업하세요.

현재 한계

• CUDA 디바이스는 하나만 지원됩니다. NVFP4 스케일 메타데이터의 샤딩 동작이 정의될 때까지 텐서 병렬 및 멀티 디바이스 device_map 구성은 거부됩니다.
• CPU 및 디스크 오프로드는 지원되지 않습니다.
• 사전 양자화된 NVFP4 체크포인트는 지원되지 않습니다.
• NVFP4 모델은 현재 save_pretrained()로 직렬화하거나 학습할 수 없습니다.