Experts4bit: fused MoE weights in 4-bit NF4/FP4#1984
Conversation
Implements per-expert block-wise 4-bit quantization for fused MoE layers (gate_up_proj + down_proj). Each expert is quantized, stored, and dequantized independently during the forward pass. - Experts4bit class in bitsandbytes/nn/modules.py - from_float() classmethod for easy construction from fp16/bf16/fp32 - Per-expert dequant loop with SiLU gated activation support - Standard state_dict serialization - Exported via bitsandbytes.nn.Experts4bit - Tests: quant round-trip, shape validation, error cases, forward vs reference, state_dict round-trip (12 tests, all pass) - Benchmarks: memory, accuracy, throughput on RTX 3090
|
Hi, This duplicates #1965, which has been open since June and implements the same Experts4bit feature. When an open PR for something already exists, the expectation is to engage with it: review, comment, suggest changes, etc; not to open a competing implementation, and certainly not one whose description proposes to close the existing PR. There was no call for competing implementations here and no prior discussion. Closing as a duplicate. If you have specific improvements over #1965, raise them there or in #1849. |
|
Hey @matthewdouglas, @pjordanandrsn — I owe you both an apology. I'm the one who opened #1984. My intention was never to compete or create confusion. I saw issue #1849 and started building my own implementation without checking whether there were already open PRs. That's completely on me. I also named it Experts4bit out of motivation. On top of that, including "Closes #1965" in my PR description made it look far more confrontational than I ever intended. I understand how that came across, and its my bad. @pjordanandrsn — your work is clearly thorough, and you got there first. I am nobody to take away from that. If anything from my implementation is useful, I'd be happy to contribute it to your PR. I have RTX 3090 benchmarks showing ~72% memory reduction and ~5000 tokens/sec throughput, along with a slightly different SiLU gating approach. If any of that helps, I'm happy to share it or help with testing and review. I'll be more careful going forward and make sure I coordinate before starting work on something that's already in progress. |
|
@zaid646, no apology needed, none was ever owed. #1849 sat open for months while the silent skip kept biting people; two of us independently building the fix isn't a collision, it's confirmation the gap is real and the design is the natural one. The name was always going to be Experts4bit. And "Closes #1965" is just what GitHub's template teaches — I read convention, not confrontation. So let's pool it: a 3090 is a card my matrix doesn't cover, and independent numbers from hardware I don't control are worth more to #1965 than anything I can produce myself. Run your benchmarks against the PR branch and post methodology + results on that thread — divergences especially welcome. And send the SiLU gating as a diff or gist; if it measures well and passes the forward-exactness suite, it lands with your name on the commit. Glad you're here. The wall's big enough for every shovel. |
Summary
Adds
bitsandbytes.nn.Experts4bit— a module for storing fused Mixture-of-Experts weights (gate_up_proj + down_proj) in 4-bit NF4/FP4 precision. Each expert is quantized independently with per-expert absmax, and dequantized on-the-fly during the forward pass, keeping the runtime memory footprint small.Closes #1965
Motivation
transformersv5 stores MoE expert weights as 3D tensors (e.g.[num\_experts, intermediate, hidden]).Linear4bitcannot wrap these because it expects 2Dnn.Linearweights.Experts4bitfills this gap.Design
nn.Parameterfor packed uint8 weightsfrom_float()classmethod for easy constructionstate_dictserializationBenchmark Results (RTX 3090)
Memory Savings
4-bit uses ~28% of fp16 memory (72% reduction, near the 75% theoretical max).
Quantization Error (NF4)
Throughput
Testing
12 tests covering quant round-trip, shape validation, error cases, forward vs reference, and state_dict round-trip. All pass. No regressions in existing test suites.