Hugging Face Transformers 怎样为推理启用半精度

用 Hugging Face Transformers 做推理，半精度不是简单地在模型后面调用一次 half()。更稳妥的做法是：加载模型时通过 dtype=torch.float16 或 torch.bfloat16 控制权重实例化；需要混合精度前向时，再用 torch.autocast 让适合降精度的算子使用半精度，其余敏感算子保留更高精度。这样既能压低显存，也能给不兼容算子留下回退边界。

先区分加载精度和计算精度

我在处理显存紧张的模型服务时，第一步不会直接把所有输入转成半精度，而是先看模型是以什么精度加载的。Transformers 当前文档把 dtype 作为 from_pretrained() 的模型权重实例化参数；dtype="auto" 则尝试读取配置中的 dtype 或旧的 torch_dtype，缺失时再参考检查点中的首个浮点权重。

这一步决定了常驻权重的显存占用。torch.autocast 是另一层机制，它只包住前向计算，让适合低精度的矩阵运算降精度，同时允许需要动态范围的算子按策略保留 float32。所以，想减少模型常驻显存时设置 dtype；想在已有模型上做混合精度前向时考虑 autocast。

根据设备选择 bfloat16、float16 或 float32

bfloat16 的指数范围接近 float32，在支持它的 CUDA 设备上通常比 float16 更容易保留数值范围；但具体吞吐仍取决于硬件、模型架构和内核。float16 兼容面更广，却更容易遇到溢出或 NaN。CPU 推理不要为了“半精度”强行改权重，先以 float32 跑通，再根据 PyTorch 与目标硬件的 AMP 支持评估 BF16。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "google/gemma-3-1b-it"  # 示例模型 ID，换成你有权限使用的检查点
device = "cuda" if torch.cuda.is_available() else "cpu"

if device == "cuda":
    # 先判断硬件是否支持 BF16，避免在不支持的设备上盲选
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
else:
    # CPU 先保持 FP32，避免把不兼容的半精度内核带进服务
    dtype = torch.float32

load_args = {"dtype": dtype}
if device == "cuda":
    # 让 Accelerate 按可用设备安排大模型权重
    load_args["device_map"] = "auto"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, **load_args)
model.eval()  # 关闭 Dropout 等训练态行为

inputs = tokenizer("用一句话解释半精度推理", return_tensors="pt")
inputs = inputs.to(model.device)  # 输入与模型保持同一设备
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=48)
print(tokenizer.decode(output[0], skip_special_tokens=True))  # 仅展示解码结果

这里的关键不是模型名称，而是加载边界：dtype 直接传给 from_pretrained()，避免先以 FP32 建立一份完整权重再转换。若模型配置已经可靠记录权重精度，也可以把它改成 dtype="auto"；这表示沿用检查点信息，不等于自动判断“哪种精度一定最快”。

用 autocast 处理混合精度前向

当模型仍以 FP32 加载，或者自定义模块对权重精度有明确要求时，可以把降精度范围限制在前向区域。PyTorch 的推荐入口是 torch.autocast(device_type="cuda", dtype=...)；推理场景不需要训练用的梯度缩放器。若某个自定义算子只接受 FP32，应把它放在禁用 autocast 的局部区域，而不是把整个模型退回 FP32。

# model 已经加载，inputs 也已放到 CUDA
amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
with torch.inference_mode():
    # autocast 只影响前向中的可转换浮点算子
    with torch.autocast(device_type="cuda", dtype=amp_dtype):
        outputs = model(**inputs)

    # 对必须保持 FP32 的后处理显式转换，避免数值边界不清
    scores = outputs.logits.float().softmax(dim=-1)
    best_token = scores[:, -1, :].argmax(dim=-1)  # 取最后位置的最高分 token

上线前用四项检查收住风险

我的判断是：显存主要卡在权重常驻时，优先改 from_pretrained(dtype=...)；显存尚可但想让部分前向算子降精度时，再引入 autocast。两者都要在目标设备上做小批量回归，不能用一次成功生成就证明半精度适合所有输入。

相关问题

可以直接对 Transformers 模型调用 model.half() 吗？

可以作为明确的实验手段，但它会修改已加载模型的参数，且可能把不适合 FP16 的模块一起转换。更可控的做法是加载时传入 dtype，或只在前向区域使用 autocast。

dtype="auto" 是否等于自动选择最佳推理精度？

不是。它主要沿用配置或检查点记录的浮点 dtype，解决加载路径和内存模式问题，不替你比较硬件吞吐、准确率和业务样本。

为什么用了半精度仍然 OOM？

权重只是显存的一部分，KV cache、激活、输入序列、批大小和 device map 都会产生额外占用。先区分加载峰值与生成峰值，再调整序列长度或批处理策略。