NVFP4

La quantification NVFP4 regroupe les poids linéaires en précision complète dans le format flottant 4 bits de NVIDIA pendant le chargement du modèle. NVFP4Config remplace les modules torch.nn.Linear éligibles sans biais, dont in_features et out_features sont tous deux divisibles par 16, par une implémentation linéaire NVFP4 issue du NVFP4 Hub kernel. Les interfaces d’attention et de MLP du modèle ne sont pas remplacées.

NVFP4 nécessite un GPU Blackwell avec une capacité de calcul 10.0 ou supérieure, une version de PyTorch compatible avec CUDA et le package kernels.

Installez Accelerate et une version compatible de kernels.

pip install --upgrade accelerate kernels

Passez NVFP4Config à from_pretrained() avec un seul appareil CUDA. Les poids sont quantifiés au moment de leur chargement, le point de contrôle source doit donc contenir des poids en virgule flottante.

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

Utilisez modules_to_not_convert pour conserver les modules sélectionnés dans leur précision d'origine.

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

Les modules linéaires NVFP4 prennent en charge torch.compile. La première invocation compilée inclut le temps de compilation du graphe, il est donc recommandé de faire chauffer le modèle avant de mesurer le débit de génération.

Limites actuelles

• Seul un appareil CUDA est pris en charge. Les configurations de parallélisme de tenseurs et de device_map multi-appareils sont rejetées tant que le comportement de fractionnement des métadonnées d'échelle NVFP4 n'est pas défini.
• Le déchargement sur CPU et disque n'est pas pris en charge.
• Les points de contrôle NVFP4 pré-quantifiés ne sont pas pris en charge.
• Les modèles NVFP4 ne peuvent actuellement pas être sérialisés avec save_pretrained() ni entraînés.