Failure running Apple MLX lora.py on 13B llms

Question

I am attempting to train my models on my own data. If I run the lora.py command using the 7b models from mistral or llama, it works fine. As soon as I upgrade to the larger models, I get an error 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}.")

I am using the Apple mlx python API on an M3 Max with 128GB RAM.

I get the exact same error running convert.py. The commands I have tried are: 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

Same error both times. I’ll post the full error below but it’s ugly.

% 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 model.layers.0.
self_attn.rotary_emb.inv_freq model.layers.9.
self_attn.rotary_emb.inv_freq model.layers.36.
self_attn.rotary_emb.inv_freq model.layers.5.
self_attn.rotary_emb.inv_freq model.layers.18.
self_attn.rotary_emb.inv_freq model.layers.6.
self_attn.rotary_emb.inv_freq model.layers.31.
self_attn.rotary_emb.inv_freq model.layers.13.
self_attn.rotary_emb.inv_freq model.layers.30.
self_attn.rotary_emb.inv_freq model.layers.29.
self_attn.rotary_emb.inv_freq model.layers.23.
self_attn.rotary_emb.inv_freq model.layers.8.
self_attn.rotary_emb.inv_freq model.layers.25.
self_attn.rotary_emb.inv_freq model.layers.20.
self_attn.rotary_emb.inv_freq model.layers.15.
self_attn.rotary_emb.inv_freq model.layers.28.
self_attn.rotary_emb.inv_freq model.layers.19.
self_attn.rotary_emb.inv_freq model.layers.12.
self_attn.rotary_emb.inv_freq model.layers.32.
self_attn.rotary_emb.inv_freq model.layers.7.
self_attn.rotary_emb.inv_freq model.layers.34.
self_attn.rotary_emb.inv_freq model.layers.37.
self_attn.rotary_emb.inv_freq model.layers.26.
self_attn.rotary_emb.inv_freq model.layers.22.
self_attn.rotary_emb.inv_freq model.layers.24.
self_attn.rotary_emb.inv_freq model.layers.16.
self_attn.rotary_emb.inv_freq model.layers.4.
self_attn.rotary_emb.inv_freq model.layers.2.
self_attn.rotary_emb.inv_freq model.layers.39.
self_attn.rotary_emb.inv_freq model.layers.11.
self_attn.rotary_emb.inv_freq model.layers.21.
self_attn.rotary_emb.inv_freq model.layers.27.
self_attn.rotary_emb.inv_freq model.layers.3.
self_attn.rotary_emb.inv_freq model.layers.1.
self_attn.rotary_emb.inv_freq model.layers.14.
self_attn.rotary_emb.inv_freq model.layers.17.
self_attn.rotary_emb.inv_freq model.layers.10.
self_attn.rotary_emb.inv_freq model.layers.35.
self_attn.rotary_emb.inv_freq.

I tried running the exact same commands on the 7B models (llama and Mistral) and they worked fine. The failures happens immediately (long before the model has been loaded).