This’ll probably be my last post on variable MoEs for a little while. It’s probably the most exciting development here, and puts a nice little bow on the whole story (for now, who knows when inspiration will strike?). Pruning MoEs by allowing for ragged expert sizes, inspired by REAP and Minitron for scoring, works really well! Being able to just shrink the experts, not fully delete them like in REAP seems to keep capability better for some models, especially as we prune more of the model. If you’re new here, I’ve been working on variable-sized experts for a while now, with some interesting results. Nothing mind blowing, but kinda cool and novel: healed OLMoE with 50% of the channels pruned scores 0.544 on GSM8K vs REAP’s 0.3171.

So, here’s the code for pruning (winnowing, I guess), here’s the updated MegaBlocks Variable (much faster than before!), a vLLM plugin adding support for these variable sized MoEs, and the model collection.

Winnowing

The main concept is that we can score each channel of each expert of an MoE to see how much it contributes to the residual stream (similar to Minitron), weighting it according to the router (borrowed from REAP). Then, since this fork of Megablocks allows for ragged experts, we can prune them according to a ranking of their scores, just like Minitron does. So, how do we do that?

Scoring the Experts

Figuring out which parts of the experts to delete is surprisingly simple actually, and of course is heavily inspired by REAP especially. REAP scores each expert and weights it based on the router, and deletes the whole expert. With winnow, we go one step lower and take out individual channels2:

# we do this for each layer and expert, skipping shared experts when applicable  
# a: [T, C] post-SwiGLU activations (input to down_proj)  
# g: [T] router weights actually applied to this expert on those tokens  
# W_down: [H, C]  
score = (g[:, None] * a.abs()).mean(0) * W_down.norm(dim=0)   # [C]  
# last, for each layer, flatten [E, C] -> rank all the channels, keep whatever %age we want

Just like in the writeup on variable FlexOlmo, I tried scoring both on general text and math text for OLMoE. It makes a pretty substantial difference (see below for the full results): about 9% of the surviving channels are different at keep-75, 17% at keep-50 and 22% at keep-25.

Pruning

With those scores (they get saved to a json file), we can prune out the corresponding weights in each expert, as blocks of 128 by default (128 is the block size of MegaBlocks). We do allow for full on deletion of experts, and especially with the smaller models, we lose a lot of experts, unsurprisingly.

OLMoE General Un-Healed Results

This method of pruning is pretty viable when we keep about 75% of the model around, but really needs healing when we get to pruning half or more of the model out. Here are the OLMoE results. In this table, we compare to REAP as well as uniform pruning of each expert– so for each expert in a layer, we prune out 50% of each one.

Unhealed REAP works great when we keep 75% of channels, but falls apart when we get down to 50% or more pruneage.

Keep Method GSM8K MATH500 IFEval HumanEval MBPP
— 🤗 download 0.691 0.250 0.612 0.348 0.302
75 winnow 0.644 0.216 0.606 0.305 0.278
75 uniform 0.589 0.150 0.603 0.244 0.222
75 REAP 0.605 0.194 0.564 0.299 0.260
50 winnow 0.447 0.134 0.488 0.213 0.172
50 uniform 0.365 0.084 0.464 0.091 0.106
50 REAP 0.051 0.028 0.233 0.006 0.042
25 winnow 0.021 0.050 0.275 0.000 0.054
25 uniform 0.036 0.034 0.309 0.018 0.000
25 REAP 0.011 0.016 0.109 0.000 0.000

OLMoE Math Un-Healed Results

Keep Method GSM8K MATH500 IFEval HumanEval MBPP
— 🤗 download 0.691 0.250 0.612 0.348 0.302
75 winnow 0.696 0.238 0.527 0.195 0.256
75 uniform 0.649 0.210 0.499 0.128 0.140
75 REAP 0.658 0.222 0.499 0.177 0.158
50 winnow 0.440 0.196 0.379 0.073 0.090
50 uniform 0.452 0.084 0.336 0.018 0.026
50 REAP 0.294 0.044 0.196 0.000 0.004
25 winnow 0.014 0.052 0.159 0.000 0.000
25 uniform 0.096 0.022 0.196 0.000 0.000
25 REAP 0.009 0.034 0.081 0.000 0.000

Qwen 3.6 35B-A3B Un-Healed Results

Same thing applied to the larger Qwen model; 75% kept is just a little worse than the unpruned model, and the more pruned the model gets, the more capacity it loses (duh). REAP is much more competitive here. Note that all of these use a generation budget of only 1280 tokens. The heavily pruned models tend to ramble, so the low scores are partially from the models not finishing.

Keep Method MMLU-Pro IFEval
— 🤗 download (FP8) 0.8402 0.8207
75 Winnow 0.763 0.802
75 REAP 0.763 0.804
50 Winnow 0.501 0.686
50 REAP 0.471 0.706
25 Winnow 0.081 0.264
25 REAP 0.210 0.438

This is kind of a fascinating result, I think. Winnow is not the clear winner here like it is above – REAP is really good for this model! My hunch at what’s going on here is that the inherited block size of 128 from Megablocks is too blunt an instrument for the smaller experts of Qwen, which has 256 experts of intermediate size 5123. Compare that to OLMoE, which has 64 experts of intermediate size 1024. Optimizing that block size is another show. I’ll come back here and edit this post when I do.

Healing the Models

As we can see, the keep-25 and keep-50 lose a lot of capability, so healing them is pretty important, and it goes pretty quickly. It worked so well in Minitron (and after that in my work on FlexOlmo), I stuck with logit distillation here, with the original, un-pruned model as the teacher. I tried both on- and off-policy (I kept the top 128 logits per token to keep some room on my disk) distillation, and they were pretty similar4. If you can amortize the teacher’s logits across a bunch of runs, off-policy is substantially faster, so that’s what I did for the OLMoE results below. For the Qwen results, since I was renting GPUs and am on a budget, I opted for on-policy distillation5 to jump straight to the healing. Since the keep-75 models are quite usable as is, I didn’t bother healing Qwen keep-75. The dataset here was Dolci-Instruct-RL.

Importantly, look at REAP below. It’s on par with Winnow for keep-75, which is great! At keep-50, second place is split between REAP and uniform, where we keep each expert but trim the intermediate size down to X% of the original (for example for OLMoE keep-25, that’s 1024 -> 256).

OLMoE General Healed Results

Keep Method GSM8K MATH500 IFEval HumanEval MBPP
— 🤗 download 0.691 0.250 0.612 0.348 0.302
75 winnow 0.663 0.227 0.616 0.354 0.278
75 REAP 0.662 0.223 0.602 0.335 0.273
75 uniform 0.617 0.207 0.617 0.278 0.232
50 winnow 0.544 0.183 0.559 0.297 0.231
50 REAP 0.317 0.137 0.517 0.240 0.217
50 uniform 0.470 0.137 0.537 0.169 0.117
25 winnow 0.207 0.084 0.443 0.020 0.112
25 REAP 0.091 0.045 0.320 0.047 0.107
25 uniform 0.177 0.057 0.428 0.059 0.004

OLMoE Math Healed Results

Same as above, approach wise, just using a math-only subset of the dataset. Stronger math results, very weak in non-math things.

Keep Method GSM8K
75 winnow 0.694
75 REAP 0.678
75 uniform 0.636
50 winnow 0.637
50 REAP 0.588
50 uniform 0.512
25 winnow 0.432
25 REAP 0.126
25 uniform 0.219

Qwen Healed Results

Again, since I wanted to get this done quickly without spending a ton to rent the compute, I used on-policy distillation here, and only healed the keep 50% and keep 25% models. Substantial healing! Healing the REAPed models is also for another time, so below are just winnowed models.

Keep MMLU-Pro IFEval
🤗 download (FP8) 0.8402 0.8207
50 0.5568 0.7597
25 0.2872 0.6765

Epilogue: Poolside Laguna-S-2.1

As a stretch goal, I wanted to try winnowing a larger model to fit nicely on my 3090s, so I opted for the Poolside Laguna-S-2.1 model. Here are 50% and 75% winnowed. At 50% winnowed and Int8 (fp8 isn’t supported on Ampere), it runs on my machine, but without a very large context window. No real useful benches here like Terminal-Bench, but it gets an 89.6 on HumanEval pass@1 (HumanEval is saturated for a model of this size, so don’t read too much into it. Think of that as just showing that it’s not totally busted.). Check it out! If you have the GPUs and want to benchmark it, I’d love to hear how it goes.

I’d love to hear from you! For feedback, or if you just want to get in touch to talk ball, feel free to email me or DM me on twitter.

  1. Cherry pick alert! Winnow does not win everywhere vs REAP. ↩

  2. It’s helpful to me to think about what’s going on with the weight matrices here. We are deleting the i-th row of the up and gate projections (these models are SwiGLU models), and the i-th column of the down projection matrix. ↩

  3. REAP is already pretty granular with this many small experts. ↩

  4. I did a not super rigorous test with off-policy distillation as the warmup and then on-policy to finish, and it raised the evals by a bit. ↩

  5. I also kinda just wanted to use OPD somewhere here, it’s cool and so hot right now. ↩

Updated: