Transformers documentation

NVFP4

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
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