Transformers documentation

NVFP4

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

NVFP4

NVFP4 量化在模型加载时,将全精度线性权重打包为 NVIDIA 的 4 位浮点格式。NVFP4Config 会将符合条件的、无偏置的 torch.nn.Linear 模块替换为来自 NVFP4 Hub kernel 的 NVFP4 线性实现,这些模块的 in_features 和 out_features 均可被 16 整除。模型的注意力机制和 MLP 接口不会被替换。

NVFP4 需要计算能力为 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 线性模块支持 torch.compile。首次编译调用会包含图编译时间,因此在测量生成吞吐量之前,请先对模型进行预热。

当前限制

  • 目前仅支持单个 CUDA 设备。在 NVFP4 缩放元数据的分片行为定义明确之前,张量并行以及多设备 device_map 配置将被拒绝。
  • 不支持 CPU 和磁盘卸载。
  • 不支持预量化的 NVFP4 检查点。
  • 目前,NVFP4 模型无法使用 save_pretrained() 进行序列化,也无法进行训练。
Update on GitHub