如何修改HuggingFace Llama模型最后一层位宽以解决推理报错？

当你用load_in_8bit加载Llama-7B模型后出现RuntimeError: "log_softmax_lastdim_kernel_impl" not implemented for 'Half'错误，核心原因是模型输出层（lm_head）被默认转为float16（Half）类型，但当前CUDA环境不支持该类型的log_softmax核函数。通过调整lm_head的精度即可解决，具体步骤如下：

1. 加载模型后修改输出层精度 #

加载模型完成后，将lm_head转换为float32（或硬件支持的bfloat16）：

import torch
from transformers import LlamaForCausalLM, LlamaTokenizer

# 8-bit方式加载模型
model = LlamaForCausalLM.from_pretrained(
    "decapoda-research/llama-7b-hf",
    load_in_8bit=True,
    device_map="auto"
)
tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")

# 将输出层转换为float32（若硬件支持bfloat16，可替换为torch.bfloat16()）
model.lm_head = model.lm_head.to(torch.float32)

2. 推理时匹配张量精度 #

推理时确保输入张量的精度与lm_head一致，避免类型不兼容：

prompt = "Hello, my name is"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# 对齐输入张量与输出层的精度
inputs = inputs.to(model.lm_head.dtype)

# 执行推理
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

补充说明 #

• 如果你的硬件支持bfloat16（如NVIDIA A100、AMD MI25等），使用torch.bfloat16()替换torch.float32()可以在精度损失极小的情况下，保持更快的推理速度。
• 该修改仅针对输出层，模型主体仍保持8-bit，不会大幅增加显存占用。

内容的提问来源于stack exchange，提问作者Lucas Azevedo