Experts4bit: fused MoE weights in 4-bit NF4/FP4#1984
Open
ZAID646 wants to merge 4 commits into
Open
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.