先把当时的情况摆出来。配置里把 bnb_4bit_compute_dtype 设成 torch.float16,加载时报 dtype 不支持,改成 bfloat16 后却能跑。

真正麻烦的地方在于,几个原因的现象几乎一样。部分模型权重本身是 bf16,与 fp16 计算类型冲突。另一种情况是显卡不支持 bf16,强行设置同样会失败。

不容易想到的是,量化存储类型与计算类型配错,矩阵乘法无法对齐。还有一种情况,不同 bitsandbytes 版本对类型组合的支持不同。

先确认显卡是否支持 bf16

python -c "import torch;print(torch.cuda.is_bf16_supported())"

然后按硬件能力设置计算类型

bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

老卡改用 fp16 并确认能加载

bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)

等前面都确认过,再打印模型首层权重类型做核对

print(next(model.parameters()).dtype)

我习惯把每一步的执行时间和结果记在一边,后面复盘省事。

恢复的判定其实很直接。模型加载成功并打印参数量,训练时前向与反向都能正常执行,显存占用保持稳定。

计算类型和存储类型是两回事,4bit 存储下用 bf16 计算得到的数值稳定性通常更好,前提是硬件支持。同一份配置在测试环境跑通,不代表线上没事,两边的默认值经常不一样。

作者 9823189