Visualizing expert firing frequencies in Mixtral MoE(mixtral-moe-vis-d726c4a10ef5.herokuapp.com)
mixtral-moe-vis-d726c4a10ef5.herokuapp.com
Visualizing expert firing frequencies in Mixtral MoE
https://mixtral-moe-vis-d726c4a10ef5.herokuapp.com/
23 comments
How does Mixtral decide which Expert is involved how much/when?
Before each feed forward block, there is a single linear layer outputting 8 numbers that tell you how promising each 'expert' is
So basically there's a 'consultant' layer.
I do wonder, say one has 16 gb of vram and 64gb of ram. IIRC mixtral takes ~ 40gb total, if you were using it for creative writing or whatever could you load certain layers into gpu memory and certain layers into system ram? I wonder what the overall speed of this would be?
You can partially offload with some backends (e.g. llama.cpp and derivatives) but speed gains from that don't come in until it's mostly offloaded. I have 8GB VRAM and it's not enough to get any boost on mixtral in Q8. 16GB might do better or it might not.
The speed is quite good even on CPU only though, I get 3.5 tokens per second with 6 cores and DDR5-6000. For comparison llama2-70B is less than 1 t/s on the same hardware in Q4. And, subjectively, Mixtral performs better.
The speed is quite good even on CPU only though, I get 3.5 tokens per second with 6 cores and DDR5-6000. For comparison llama2-70B is less than 1 t/s on the same hardware in Q4. And, subjectively, Mixtral performs better.
I did this with 24GB VRAM and 32GB DDR5, using LM Studio, and it was about as fast as I could read. (I read fast but I'd have to run it again to guess the token rate.)
I'm upgrading to 96GB RAM now to run the larger models, but I do wonder whether it'll be slow when using proportionally less VRAM.
I'm upgrading to 96GB RAM now to run the larger models, but I do wonder whether it'll be slow when using proportionally less VRAM.
If you are making many requests in batch this works ok because you can shuffle the next layer in while the current one is processing a set of matrix multiplies. This takes it from being a memory bound problem to a flops bound problem. This really only works if you care about throughput and not latency.
I understand that for each token mixtral will only need two (of eight) submodels. I wonder if there is temporal locality and an LRU caching schema could be used.
It is two out of eight at each layer, with 32 layers independent of each other. There are no eight "sub-models".
However, this raises a question: could a slightly more complex router use output layer n-1 to choose experts for layer n+1 (vs n and n+1 today)? This way, there is more time to load the needed experts for the n+1 layer.
However, this raises a question: could a slightly more complex router use output layer n-1 to choose experts for layer n+1 (vs n and n+1 today)? This way, there is more time to load the needed experts for the n+1 layer.
With a recent llama.cpp, Mixtral, when run using CPU only (8 threads), produces 4 tokens per second on a Dell Inspiron 7415 laptop with AMD Ryzen 7 5700U, upgraded to 64 GB of RAM. As this is an APU, even one layer is too much for the VRAM, so acceleration cannot be used.
Does anyone know of any existing research into the effects of scaling the number of experts up in models like this? I assume there's diminishing returns and an eventual inversion of quality (maybe?), but it'd be interesting to see the curve.
There is an 'old' paper https://arxiv.org/abs/1701.06538 but I believe there needs to be renewed effort in this direction.
I don't understand why there are 31 layers in this tool.
I don't understand why we only see 4 charts (I would expect to see 8, one for each 7B mistral tiny model)
I don't understand why we only see 4 charts (I would expect to see 8, one for each 7B mistral tiny model)
The plots show 2 dimensional projection of the 8 dimensional feature vector of each paragraph. So, x and y axis are linear combination of 8 different experts. Ideally, all of this should be in a single plot but there was a lot of overlap between different sub-categories and it was hard to see. So, I thought of separating them by their broad categories. Also, there are 32 layers in the model, each has 8 different experts, and 2 of them picked for each token.
Can't visualize the graphs well from mobile but the many layers are there because MoE models select a different expert at each layer, not just at each token.
There are 32 layers (numbered from 0), which is... almost how many the model has. Actually it's got 33, but maybe one of them is before or after the separate experts?
The base model has 32 layers and there is a single linear layer for language modeling (going from embeddings to the vocabulary) that gets applied at the very end.
Interesting implications for future optimization. If you don't have enough RAM to fine-tune the whole model, you could use this to select the two most relevant experts for each layer of the model. Maybe for inference, you could quantize the less-relevant experts to be smaller.
Fixing the experts for a layer might not work since all experts fire almost with equal probability. There are small variations by topic but they are consistent enough to be captured with a simple linear classifier. I believe this happens due to the load balancer loss which forces the model to pick all experts with equal probability. However, what you're saying is a great direction for future MoE's. Can we train MoEs without load balancing so that it is possible to quantize/prune the non-relevant experts more aggressively? We haven't had any major open source MoE's because, as far as I know, they are not straightforward to train, but I expect this to change.
I could be reading this wrong, but doesn't the article say that the MMLU topic is enough to determine which experts got picked 96% of the time?
Edit: Oh I read it completely backwards, didn't I. Given the decisions, you can determine the topic of the input.
Edit: Oh I read it completely backwards, didn't I. Given the decisions, you can determine the topic of the input.
I would assume that you would get a big performance drop, because knowledge is split in a very non human understandable way between the experts
So use a non human algorithm to optimize that problem by maybe collecting some huristics during a benchmark phase, and then the quantization step can take that into account based on specific pathways that are never hit because another expert is always better at that task and always gets picked. It seems like there is a ton of room for improvement there.
Does anyone have the resources to try this with PHATE (https://www.nature.com/articles/s41587-019-0336-3)