如何解决TypeError：dispatch_model()收到意外关键字参数'offload_index'？

问题场景 #

在Kaggle环境中尝试将Hugging Face的facebook/opt-13b大模型拆分权重加载到多GPU时，运行指定代码触发上述TypeError。

错误原因 #

这是transformers与accelerate库版本不兼容导致的：新版本transformers调用dispatch_model函数时会传入offload_index参数，但当前环境的accelerate版本较低，该版本的dispatch_model并未实现这个参数，因此触发参数不匹配错误。

解决方法 #

方法一：升级accelerate到兼容版本

在Kaggle代码单元格中执行以下命令升级accelerate：

!pip install --upgrade accelerate

升级完成后重启内核，再重新运行模型加载代码即可。

方法二：使用accelerate官方推荐的加载方式（更稳妥）

替换原加载代码为load_checkpoint_and_dispatch函数，这是accelerate专门为大模型多设备拆分/Offload设计的加载方法，版本兼容性更好：

import torch
from accelerate import infer_auto_device_map, init_empty_weights, load_checkpoint_and_dispatch
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")

# 改用load_checkpoint_and_dispatch加载模型
model = load_checkpoint_and_dispatch(
    model,
    checkpoint="facebook/opt-13b",
    device_map=device_map,
    offload_folder="offload",
    dtype=torch.float16
)

额外注意 #

如果升级库后出现依赖冲突，可尝试添加--force-reinstall参数重新安装，同时确保transformers版本不低于4.20.0（可通过!pip install --upgrade transformers同步升级）。

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