NVFP4

La cuantización NVFP4 empaqueta los pesos lineales de precisión completa en el formato de punto flotante de 4 bits de NVIDIA mientras se carga el modelo. NVFP4Config reemplaza los módulos torch.nn.Linear elegibles sin sesgo, cuyos in_features y out_features son ambos divisibles entre 16, con una implementación lineal NVFP4 del NVFP4 Hub kernel. Las interfaces de atención y MLP del modelo no se reemplazan.

NVFP4 requiere una GPU Blackwell con capacidad de cómputo 10.0 o superior, una compilación de PyTorch habilitada para CUDA y compatible, y el paquete kernels.

Instala Accelerate y una versión compatible de kernels.

pip install --upgrade accelerate kernels

Pasa NVFP4Config a from_pretrained() con un único dispositivo CUDA. Los pesos se cuantizan a medida que se cargan, por lo que el punto de control de origen debe contener pesos en coma flotante.

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

Use modules_to_not_convert para mantener los módulos seleccionados en su precisión original.

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

Los módulos lineales NVFP4 admiten torch.compile. La primera invocación compilada incluye el tiempo de compilación del grafo, por lo que debe calentar el modelo antes de medir el rendimiento de generación.

Limitaciones actuales

• Solo se admite un dispositivo CUDA. Las configuraciones de paralelismo de tensores y de device_map para varios dispositivos se rechazan hasta que se defina el comportamiento de particionado de los metadatos de escala NVFP4.
• No se admite la descarga a CPU ni a disco.
• No se admiten puntos de control NVFP4 pre-cuantizados.
• Los modelos NVFP4 no pueden serializarse actualmente con save_pretrained() ni entrenarse.