-
Notifications
You must be signed in to change notification settings - Fork 823
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Enable the Mixtral-like Moe model without the quantized gate layer #394
Comments
Actually, if the size is not 8 won't the gating layer just be quantized? I feel like it would be better to simply quantize the gate layers so we don't have to deal with mixed quantization. As far as I understand the only reason we have the predicate for |
Yeah, I had thought about just the quantize gate. but, the gate matrix is quite small (4096*(3-8)). From my experiment with the current mixtral implementation, I didn't see much effect of applying the lora weight to it (maybe there's an issue with the current gate implementation in mixtral?). However, when I tried fine-tuning the gate layer directly, I encountered a loss value of nan, I am not sure if lora on quant gate works or not. |
@awni Would you be able to give me some hints as to why we have to use stop_gradient indices on the sparse MOE block? Also, is there anything I can do to help fine-tune the Moe routing and make it work properly? Personally, I believe that there is great potential for Moe models to excel on Apple Silicon. Thanks in advance. |
Well it's a good question. We could have a gradient on the Training MOEs in general requires a bit of finesse. It's worth taking a look at the literature to see what kind of techniques they use to balance expert selection. Usually there is some forced hard constraint like each expert has to be used a certain number of times per token.
I don't know exactly what you mean by "work properly". I think we need to understand what exactly is not working and diagnose the reason. Is the current behavior that the gates are collapsing to only one or two experts? I would be happy to dig into this a bit when I have some bandwidth, possibly early next week. In the meantime we can use this thread to chat about it? And if you uncover anything interesting I would definitely love to hear about it / give my input. |
@awni, didn't expect you guys to spend too much effort on this right now. I know you guys are busy improving the mlx framework 🚀 , but just a few hints or thoughts that could help me investigate the issue would be greatly appreciated. Here are some of my findings. I am keen to hear your thoughts or if there is anything I missed.
|
Just found that I removed the |
Btw, there were some issues with eval'ing in grad graphs in introduced in Are you using the latest MLX for these tests? If so I might go back to 0.0.11 until we do a patch release or you can use main once my fix PR lands. |
Double checked with mlx latest master, it seems that the backpropagation still doesn't work with argpartition. I have to add the stop_gradient back. I don't know if there is a stop_gradient on argpartition. Will the fine-tuning with expert's MLP still work? I couldn't wrap my head around how the backpropagation flows through the topk in this case. |
I think to get backprop to work through expert index selection we would have to do something like:
|
@mzbac a couple requests from you:
|
@awni I am using mlx-moe (https://github.com/mzbac/mlx-moe) as a playground to test moe fine-tuning. It is easy to set up the training by following the readme. However, it defaults to using mixtral 4x7b. Tomorrow, I will see if I can create a phi-2 4x moe model for quick testing. Or, if you prefer mlx-lm command, you can use the following command:
Regarding the fine-tuning part, I noticed that the gate weights are being updated, but the changes in values seem very small. I didn't analyze the results thoroughly because during training, I felt that there wasn't much decrease in loss. My guess was that maybe only fine-tuning q,v and gate is not sufficient. However, tomorrow I will do more testing on my local machine to see if I can find anything. Thoughts: |
@awni, I have created a phi2 2x4 moe model with 4 identical phi2 experts. You can check out my fork branch here: https://github.com/mzbac/mlx-examples/tree/mlx-lm/moe. (I need to disable quant to gate, somehow there is an issue with lora the quant gate, I getting loss nan issue, also need to add phixtral model support) python -m mlx_lm.lora --model mzbac/phi2-2x4-hf-4bit-mlx --train --iters 600 --data ../lora/data --lora-layers 32 I have tested it on my local machine and it seems to work well with q,v, and gate lora. However, I assume that this is because the wikisql training dataset has very simple patterns to learn from, so even a regular phi2 model can be easily fine-tuned without much help from the experts. This means that the gate routing doesn't really matter in this case. In my experiment of fine-tuning on a complex dataset (guanaco) without being able to lora the experts' mlp layers, the loss doesn't decrease much (always around 1.xx). |
Could you point me to this dataset? I can try training on it to see if I can repro
So you think there is a capacity issue? Did you try training a larger (but non-MOE) model to see if it can learn this dataset? It might not be related to the experts...could be some other setting in the training config /setup. |
Yeah, just try
And also, I tried on a guanaco here. the loss seems always around 0.9x-1.xx. |
Sorry I am still not understanding 100%:
This is with WikiSQL right? Just curious, where did the loss end up?
What's stopping you from doing this (other than changing the default settings). Does it crash if you make the MLP layers LoRA layers? |
Regarding capacity issues, some simple things to try are:
If those two tacks are not effective then it suggests something is wrong with the optimization.. |
I tried around 2000 iterations (around 3 epochs). The loss is approximately 0.6 or 0.5x, but when I ran the inference, the result was not very good.
Maybe I misunderstood. If the moe doesn't backpropagate at topk, can we still apply LORA to the expert's MLP layers? I thought it wouldn't be able to update LORA since the backprops stop at argpartition inds. |
@awni please let me know if raise the issue in mlx for support torch topk like ops would be making more sense for this? |
I think it’s a good idea to have issues in mlx for the missing grads. I’m still not certain that is the problem here though. Sorry I have been intending to look into this just after we get our next release out. What is a good dataset to use to investigate? Sounds like wikisql is a bit too on the easy side? |
I think there are two questions, I am not very sure.
For dataset I feel wikisql may too easy to overfit so it doesn't require lora mlp layers. But I think if we enabled full fine-tuning all linear layer then it will be up to user to adjust the hpye params than mlx concerns. Please let me know if those make sense, I am happy to setup some playground repo for testing. |
Yes there can still be gradient to those layers. Did you try adding LoRA to the MLP layers? You can check if the gradient is zero during the update but I do not think it should be.
The conversion to NumPy there is not ideal .. but I don't think it should preclude one from training the expert MLP. There should still be gradient |
@awni, thanks for clarifying it. I will do some tests. I haven't really checked the weight of the mlp's lora before, was assuming it wouldn't work. |
Awesome, let me know how it goes!! I was just planning to try some MOE fine-tunes to see how it all works myself. |
@mzbac is that dataset a good one to try to see how well the MOE fine-tuning works? |
Yeah, Not being experts on this topic, but that's the dataset I have used to fine-tune a custom moe models. From my experience, it can reduce loss well during the fine-tuning process and doesn't easily lead to overfitting (I have tried 9k-15k iterations and didn't see overfitting on phi2-2x3 and qwen2-2x3). However, if you want to try something with fewer examples, I personally think https://huggingface.co/datasets/timdettmers/openassistant-guanaco also would be a good option. The guanaco is not in pair format, so preprocessing may be required if you intend to fine-tune chat-based models. |
@mzbac I tried LoRA fine tuning Phixtral on the wizard dataset and it works pretty well as far as I can tell:
Is that what you expect? Maybe that dataset is also too easy? This is the setup I used:
I'm guessing the gates did not change much though but I didn't check yet. |
Very strange, I have never get a loss below 0.5 on the wizard dataset before, but I did apply Lora with 32 layers of q, v and gate. Do you think that having all those layers might be causing some instability? I can try the 16 layers see if I can reproduce the results. |
@awni, one thing: the Phixtral doesn't have a noise linear layer in the implementation. So during fine-tuning, the gate may quickly tend to direct all tokens to a favorite expert. My moe implementation has this noise layer enabled during fine-tuning https://github.com/mzbac/mlx-moe/blob/phi2-moe/phi2moe.py#L170-L179 |
Right I did not train with that. It might be collapsing to one expert.. I will see if it works with the noise |
So the model (as you've pointed out before) starts out only using two experts for every token, and it stays that way during Lora finetuning. I checked the grad of the gate and it is zero 🤔
|
Yeah, I am currently running llm_evaluation for the qwen 2x3 model. It actually performs worse in benchmarks compared to the original 7b chat model. I suspect this might be related to it. MMLU
GSM8K
Qwen1.5-7B-Chat
GSM8K
|
@mzbac there was a bug in my code so disregard the zero gradient comment above. I fixed that. SetupI did some analysis to understand how the expert selection changes over training with Phixtral. Some comments on setup:
Command:
Findings:
Layer 0 does not learn experts, only expert 1 and 2 are selected, 3 and 4 have 0 counts (lines are on top of each other): ![]() Layer 16 is the first LoRA layer, all experts get some tokens: ![]() Same for Layer 31: ![]() |
So it seems like learning the gate layer is working fine 🤔 ... I'm wondering what we are doing differently? I added the changes to work with Phixtral to this PR #426 (and removed phixtral as a standalone model from |
@awni Thanks for the thorough analysis. I will try your setup and see if I can reproduce the result. I will report back if I find anything. |
@awni I have done some testing and here are my findings: Test 1: Test 2: Test 3: I can confirm that the gate is definitely getting updated and introducing noise into it might be necessary in order to avoid imbalanced routing. |
Thanks for sharing.
Interesting. Are you measuring the loss with or without the noise addition? I assume it is better to measure without noise addition since you would not use it when actually generating with the model. Maybe that accounts for the discrepancy? |
Just sharing the findings, I have done a full linear layers Lora fine-tuning on one phi2-2x4 model. It looks like the experts' MLP is updating well despite the noise in the linear layer. even, the loss does not seem to decrease and remains around 0.6x. Nevertheless, the MMLU evaluation indicates an improvement in MMLU score. I will try to do some fine-tuning without the noise layer later this week and see if there are any differences. |
That’s super cool! Thanks for the update! |
@mzbac Do you mind sharing the harness you used for the MMLU benchmark? I'd like to run evaluations on my MLX models but haven't done anything like this before |
Yeah, I am using https://github.com/EleutherAI/lm-evaluation-harness. It's slow on the MPS backend, so you have to run it on a CUDA device. I normally run the benchmarks on my PC instead of mac |
Thank you. I was poking around that very tool, how it evaluates, and if it can be done over Open AI to calculate perplexity of MMLU tasks (topics) on mlx fine tuned models in the same way the test/evaluate libraries do. (Sorry if this is off topic, maybe later I can add a discussion item) |
I only used basic evaluation and didn't set up evaluation with OpenAI. I think this is quite a big topic; maybe we can have a separate discussion. In my opinion, currently the LLM evaluation is kind of lagging behind the progress of the LLM advancements, so I normally just check the benchmarks difference over different epochs without really paying much attention to the actual score. |
Currently, the community has started experimenting with building more models using a mix of different local experts. In the current implementation of mlx-lm, we have hardcoded the linear_class_predicate with 8 as an assumption that there will always be 8 local experts. It would be great if we could read the number of local experts from the configuration and make it adjustable to support varying numbers of local experts. This way, users can use mlx-lm to fine-tune the gate for LoRa. For example, here I have to copy and paste code from mlx-lm in order to make lora moe 4x7b work.
The text was updated successfully, but these errors were encountered: