在Apple MLX中运行lora.py训练13B大模型时权重加载失败

问题描述 #

在配备128GB内存的M3 Max设备上使用Apple MLX Python API训练自有数据：

• 运行Mistral/Llama 7B模型时，convert.py和lora.py命令可正常执行
• 切换至13B模型（如meta-llama/Llama-2-13b-hf）或Mixtral-8x7B-Instruct-v0.1时，模型加载阶段直接触发报错：

  ValueError: Received parameters not in model: model.layers.*.self_attn.rotary_emb.inv_freq
  

  报错源自mlx/nn/layers/base.py文件的load_weights方法第211行。

执行命令示例：

python convert.py --hf-path mistralai/Mixtral-8x7B-Instruct-v0.1 -q
python lora.py --model mistralai/Mixtral-8x7B-Instruct-v0.1 --train --batch-size 1 --lora-layers 4

完整报错信息：

% python lora.py --model meta-llama/Llama-2-13b-hf --train --batch-size 1 --lora-layers 4

Loading pretrained model

Fetching 11 files: 100%|...| 11/11 [00:00<00:00, 115922.97it/s]
Traceback (most recent call last):
File "/Users/conleysa/AI/MLX/mlx-examples/lora/lora.py", line 321, in
model, tokenizer, _ = lora_utils.load(args.model)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/conleysa/AI/MLX/mlx-examples/lora/utils.py", line 140, in load
model.load_weights(list(weights.items()))
File "/Users/conleysa/AI/MLX/mlx_venv/lib/python3.11/site-packages/mlx/nn/layers/base.py", line 211, in load_weights
raise ValueError(f"Received parameters not in model: {extras}.")
ValueError: Received parameters not in model: model.layers.38.self_attn.rotary_emb.inv_freq model.layers.33.self_attn.rotary_emb.inv_freq ...（省略重复项）

解决办法 #

1. 更新MLX及示例代码 #

MLX早期版本对大模型的Rotary Embedding参数处理存在兼容问题，执行以下命令更新依赖：

pip install --upgrade mlx mlx-lm
cd mlx-examples && git pull

2. 过滤不必要的权重参数 #

若更新后仍报错，修改mlx-examples/lora/utils.py中的load函数，在加载权重前过滤掉inv_freq相关参数：
找到model.load_weights(list(weights.items()))这一行，替换为：

# 过滤rotary_emb.inv_freq参数
filtered_weights = {k: v for k, v in weights.items() if not k.endswith("rotary_emb.inv_freq")}
model.load_weights(list(filtered_weights.items()))

3. 转换模型时跳过参数 #

使用convert.py转换模型时，添加--ignore参数跳过inv_freq：

python convert.py --hf-path mistralai/Mixtral-8x7B-Instruct-v0.1 -q --ignore "model.layers.*.self_attn.rotary_emb.inv_freq"

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