无CUDA环境使用ChatGLM-6B遇BFloat16 dtype错误求助

问题原因 #

ChatGLM-6B默认加载BF16精度的权重，在无CUDA的CPU环境下直接移除.half().cuda()后，输入张量的 dtype 与模型权重的 dtype 不匹配，触发该错误。

解决方法 #

方法一：指定模型加载为float32精度

修改模型加载代码，明确指定torch_dtype="float32"，统一CPU运行时的 dtype：

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, torch_dtype="float32")
response, history = model.chat(tokenizer, "你好", history=[])
print(response)
response, history = model.chat(tokenizer, "晚上睡不着应该怎么办", history=history)
print(response)

方法二：加载量化版本模型（推荐）

使用INT4或INT8量化版的ChatGLM-6B，既解决 dtype 冲突，又大幅降低CPU运行时的内存占用：

from transformers import AutoTokenizer, AutoModel
# 加载INT4量化版
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
response, history = model.chat(tokenizer, "你好", history=[])
print(response)
response, history = model.chat(tokenizer, "晚上睡不着应该怎么办", history=history)
print(response)

注：如果你的CPU支持BF16指令集（如较新的Intel Ice Lake或AMD Zen3及以上），也可以尝试指定torch_dtype="bfloat16"，但该方案兼容性不如前两种。

内容的提问来源于stack exchange，提问作者גיא דורנפלד