Target MLP weights in MLX-lm?

[mark-lord]

As per this discussion, I managed to get rank to 128 and QKVO layers on (yay!) but unfortunately no luck on the mlp weights yet. Here's the code I'm using for the lora utils.py tuning file in mlx-lm:

check_lora_layers(len(model.model.layers))

        for l in model.model.layers[len(model.model.layers) - num_lora_layers :]:
            l.self_attn.q_proj = LoRALinear.from_linear(l.self_attn.q_proj)
            l.self_attn.v_proj = LoRALinear.from_linear(l.self_attn.v_proj)
            l.self_attn.k_proj = LoRALinear.from_linear(l.self_attn.k_proj)
            l.self_attn.o_proj = LoRALinear.from_linear(l.self_attn.o_proj)
            '''l.mlp.gate_proj = LoRALinear.from_linear(l.mlp.gate_proj)
            l.mlp.down_proj = LoRALinear.from_linear(l.mlp.down_proj)
            l.mlp.up_proj = LoRALinear.from_linear(l.mlp.up_proj)'''
            if hasattr(l, "block_sparse_moe"):
                l.block_sparse_moe.gate = LoRALinear.from_linear(
                    l.block_sparse_moe.gate
                )

When I un-comment the lines about mlp, I get this error:

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 175, in <module>
    linear_to_lora_layers(model, args.lora_layers)
  File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx_lm/tuner/utils.py", line 43, in linear_to_lora_layers
    l.mlp.gate_proj = LoRALinear.from_linear(l.mlp.gate_proj)
                                             ^^^^^
  File "/Users/marklord/Documents/Fine-tuning/MLX/venv/lib/python3.11/site-packages/mlx/nn/layers/base.py", line 101, in __getattr__
    raise AttributeError(f"{type(self)!r} has no attribute {key!r}")
AttributeError: <class 'mlx_lm.models.mixtral.MixtralDecoderLayer'> has no attribute 'mlp'

Am I listing the mlp weights wrong?

[awni]

For Mixtral you will have to add the lora weights to each expert separately and they have different names. The code I sent above is meant for Mistral / Lama style models.

For example:

l.block_sparse_moe.experts[i].w1

Should work for the first weight of the ith expert.

[awni]

The model file has all the args etc: https://github.com/ml-explore/mlx-examples/blob/main/llms/mixtral/mixtral.py

[mark-lord]

The speed at which you reply to requests is mind-boggling, thank you so much 😄 really appreciate everything you're building here

Will test out those fixes when I get time later :)

[mzbac]

Here is an example of how I fine-tuned the Moe model for MLP layers, you can adjust it to fit your specific use case.

    model.lm_head = LoRALinear.from_linear(model.lm_head, r=64, lora_alpha=128)
    for l in model.model.layers:
        l.self_attn.q_proj = LoRALinear.from_linear(
            l.self_attn.q_proj, r=64, lora_alpha=128
        )
        l.self_attn.v_proj = LoRALinear.from_linear(
            l.self_attn.v_proj, r=64, lora_alpha=128
        )
        l.block_sparse_moe.gate = LoRALinear.from_linear(
            l.block_sparse_moe.gate, r=64, lora_alpha=128
        )

        for e in l.block_sparse_moe.experts:
            e.gate_proj = LoRALinear.from_linear(e.gate_proj, r=64, lora_alpha=128)
            e.down_proj = LoRALinear.from_linear(e.down_proj, r=64, lora_alpha=128)
            e.up_proj = LoRALinear.from_linear(e.up_proj, r=64, lora_alpha=128)