QLoRA for Mixtral?

[mark-lord]

Question is the title; is there QLoRA support for Mixtral by any chance? I set model to Mixtral-8x7B-Instruct-v0.1 and tried it out on my Jupyter Notebook but got this error:

Loading pretrained model
Traceback (most recent call last):
File "", line 198, in _run_module_as_main
File "", line 88, in _run_code
File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx_lm/lora.py", line 170, in
model, tokenizer = load(args.model)
^^^^^^^^^^^^^^^^
File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx_lm/utils.py", line 356, in load
model = load_model(model_path, lazy)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx_lm/utils.py", line 320, in load_model
model.load_weights(list(weights.items()))
File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx/nn/layers/base.py", line 164, in load_weights
raise ValueError(f"Received parameters not in model: {extras}.")
ValueError: Received parameters not in model: model.layers.31.mlp.up_proj.weight model.layers.21.mlp.gate_proj.biases model.layers.26.mlp.up_proj.biases model.layers.0.mlp.down_proj.scales model [..etc]

I assume by this that it's not supported :'(

[awni]

Can you share the command? Mixtral should work

[mark-lord]

Running in Jupyter, ran these two cells and got these two outputs. This code works for Mistral-7b

!python -m mlx_lm.convert --hf-path "mistralai/Mixtral-8x7B-Instruct-v0.1" -q

[...etc] 
Fetching 26 files: 100%|████████████████████| 26/26 [00:00<00:00, 110600.31it/s]
[INFO] Quantizing

!python -m mlx_lm.lora
--train
--model ./mlx_model
--data ./data
--batch-size 1
--lora-layers 32
--iters 100
--max-seq-length 256
--learning-rate 2e-5

Loading pretrained model
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx_lm/lora.py", line 170, in <module>
    model, tokenizer = load(args.model)
                       ^^^^^^^^^^^^^^^^
  File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx_lm/utils.py", line 356, in load
    model = load_model(model_path, lazy)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx_lm/utils.py", line 320, in load_model
    model.load_weights(list(weights.items()))
  File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx/nn/layers/base.py", line 164, in load_weights
    raise ValueError(f"Received parameters not in model: {extras}.")
ValueError: Received parameters not in model: model.layers.16.mlp.up_proj.biases model.layers.2.mlp.down_proj.scales model.layers.16.mlp.gate_proj.weight model.layers.18.mlp.down_proj.scales model.layers.0.mlp.gate_proj.weight model.layers.4.mlp.up_proj.biases [...etc] model.layers.29.mlp.gate_proj.biases model.layers.21.mlp.down_proj.biases.

[mark-lord]

Might be worth mentioning I also keep getting this warning:

!python -m mlx_lm.convert --hf-path "mistralai/Mixtral-8x7B-Instruct-v0.1" -q

<frozen runpy>:128: RuntimeWarning: 'mlx_lm.convert' found in sys.modules after import of package 'mlx_lm', but prior to execution of 'mlx_lm.convert'; this may result in unpredictable behaviour [INFO] Loading Fetching 26 files: 100%|█████████████████████| 26/26 [00:00<00:00, 15169.27it/s] [INFO] Quantizing

[mark-lord]

Naïvely thinking it may be a problem with having done Mistral-7b first and then swapped out to do Mixtral without having deleted the 7b first. Either way my MLX_Model directory looks like this:

Opening up config, it at least seems to be for Mixtral: "num_experts_per_tok": 2, "num_local_experts": 8.

[awni]

Oh that's possible. I would definitely do separate directories for each model you convert. You can specify the ouput directory to convert with --mlx-path

[mark-lord]

That fixed it! Looking at the new folder vs the old, the only difference is that model.safetensors file. Looking at that file, it's the exact same size as the Mistral-7b's one that it was using. Interestingly everything else related to Mistral-7b has been removed from that directory. So if I had to guess, I'd say that it's not meant to be there and was meant to have been cleaned up automatically? And MLX was looking at it for the training but wasn't able to find the layers since the config didn't match up to the safetensors file itself? Either way, this was the answer :)

[awni]

Yea that file was the problem. We don't do any automatic cleanup, probably we should clear the model directory before writing to it..