macOS环境下Python使用BFloat16加载Bloom模型报错问题

问题复现 #

在macOS系统上使用bigscience/bloom模型实现文本补全，使用的测试脚本如下：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom')
inputs = tokenizer("My dog is ", return_tensors="pt")

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", 
offload_folder="offload", torch_dtype=torch.bfloat16)
output = model.generate(inputs["input_ids"].to(0), min_length=30, max_length=30, 
  do_sample=True)
print(tokenizer.decode(output[0].tolist()))

本地环境为Python 3.9.12，已通过pip3安装accelerate、torch、transformers库，运行脚本时模型加载阶段抛出如下错误：

Traceback (most recent call last):
  File "/Users/stevex/bloom/download.py", line 8, in <module>
    model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", offload_folder="offload", torch_dtype=torch.bfloat16)
  File "/opt/homebrew/lib/python3.9/site-packages/transformers-4.21.0.dev0-py3.9.egg/transformers/models/auto/auto_factory.py", line 446, in from_pretrained
    return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
  File "/opt/homebrew/lib/python3.9/site-packages/transformers-4.21.0.dev0-py3.9.egg/transformers/modeling_utils.py", line 2119, in from_pretrained
    model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
  File "/opt/homebrew/lib/python3.9/site-packages/transformers-4.21.0.dev0-py3.9.egg/transformers/modeling_utils.py", line 2332, in _load_pretrained_model
    new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
  File "/opt/homebrew/lib/python3.9/site-packages/transformers-4.21.0.dev0-py3.9.egg/transformers/modeling_utils.py", line 553, in _load_state_dict_into_meta_model
    offload_index = offload_weight(param, param_name, offload_folder, offload_index)
  File "/opt/homebrew/lib/python3.9/site-packages/accelerate/utils/offload.py", line 25, in offload_weight
    array = weight.numpy()
TypeError: Got unsupported ScalarType BFloat16

尝试执行pip3 install bfloat16安装第三方bfloat16库解决问题时，编译bfloat16.cc文件失败，典型编译错误如下：

bfloat16.cc:462:26: error: unexpected type name 'bfloat16': expected expression
                              /*alignment=*/alignof(bfloat16),
                                                    ^
      bfloat16.cc:757:43: error: a space is required between consecutive right angle brackets (use '> >')
                      struct TypeDescriptor<std::complex<float>>
                                                              ^~
                                                              > >
      bfloat16.cc:764:44: error: a space is required between consecutive right angle brackets (use '> >')
                      struct TypeDescriptor<std::complex<double>>
                                                               ^~
                                                               > >

问题原因 #

这个报错和macOS系统本身支不支持BFloat16没有直接关系，核心诱因有两个：

• 安装的transformers是4.21.0开发版，配套的旧版accelerate在做权重离线卸载（offload）时，会强制把张量转成numpy格式存储，但numpy原生不支持BFloat16数据类型，直接转换就会抛出类型错误。
• 尝试安装的第三方bfloat16库长期未维护，C++代码是按GCC编译器规范编写的，和macOS默认的Apple Clang编译器语法要求不兼容，所以编译失败，就算强行编译安装也解决不了accelerate的转换逻辑问题，没必要在这个库上浪费时间。

解决步骤 #

按优先级选以下方案即可：

• 优先更换权重加载精度（最省事，兼容性最好）
  macOS环境下不管是Intel芯片还是M系列Apple Silicon芯片，跑模型优先用float16或float32精度，完全不需要硬上bfloat16。把模型加载代码里的torch_dtype=torch.bfloat16替换成对应精度即可：
  • M系列芯片Mac：用torch_dtype=torch.float16，内存占用低，运算速度快
  • Intel芯片无独立显卡的Mac：用torch_dtype=torch.float32，避免float16无硬件加速导致的运算报错
    另外注意原代码里inputs["input_ids"].to(0)是强制把输入放到0号CUDA设备（NVIDIA显卡），macOS没有NVIDIA CUDA环境会直接报错，要改成自动匹配模型所在设备。修改后的可运行代码示例：

  import torch
  from transformers import AutoTokenizer, AutoModelForCausalLM
  
  tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom-560m') # 普通Mac建议先跑560M小参数版本，全量BLOOM模型内存要求极高
  inputs = tokenizer("My dog is ", return_tensors="pt")
  
  # M系列芯片用这行
  model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", device_map="auto", 
  offload_folder="offload", torch_dtype=torch.float16)
  # Intel无独显Mac替换成下面这行
  # model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", device_map="auto", 
  # offload_folder="offload", torch_dtype=torch.float32)
  
  output = model.generate(inputs["input_ids"].to(model.device), min_length=30, max_length=30, 
    do_sample=True)
  print(tokenizer.decode(output[0].tolist()))
  

• 如果一定要用bfloat16精度
  先把相关依赖升级到最新正式版，新版transformers和accelerate已经修复了bfloat16权重卸载时的numpy转换bug，M系列芯片的PyTorch MPS后端也原生支持bfloat16运算，不需要安装第三方bfloat16库。执行升级命令：

  pip3 install --upgrade torch transformers accelerate
  

  升级完成后即可正常加载bfloat16精度的模型。

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