Add Experts4bit for 4-bit quantization of fused MoE experts

[pjordanandrsn]

Problem

transformers v5 stores MoE expert weights as a single fused 3-D nn.Parameter
([num_experts, out_features, in_features], e.g. OlmoeExperts, Qwen3MoeExperts) instead of an
nn.ModuleList of nn.Linear. bitsandbytes' 4-bit conversion replaces nn.Linear modules, so these
fused experts — the dominant share of an MoE's weights — are silently skipped and stay full precision
(#1849).

What this adds

bitsandbytes.nn.Experts4bit: a small nn.Module holding the two fused expert projections
(gate_up_proj, down_proj) in 4-bit NF4/FP4.

• Plain nn.Parameter (uint8) + absmax buffers, not Params4bit. A 3-D per-expert stack doesn't
  fit Params4bit's tensor-subclass / device-movement machinery cleanly; keeping packed bytes +
  per-expert absmax as ordinary tensors means the module serializes through the standard state_dict
  with no custom save/load hooks (safetensors.torch.save_model/load_model work out of the box —
  pinned by a test).
• Per-expert quantization (from_float loops over experts, quantize_4bit each). in_features must
  be a multiple of blocksize so blocks never straddle an expert boundary (validated).
• Forward dequantizes one expert at a time and runs the SwiGLU expert MLP, mirroring the reference
  fused-experts forward. It restores the [packed, 1] layout quantize_4bit emits so
  dequantize_4bit's transpose back-compat shim (keyed on A.shape[0] == 1) doesn't fire.
• Training memory — re-dequantize in backward instead of saving. Every expert projection routes
  through _FrozenLinearRecomputeBackward: forward is dequantize + F.linear; backward re-dequantizes
  the frozen weight to form grad_output @ W. The dequantized [out, in] weight never enters autograd's
  saved-tensor storage, so training activation memory stays independent of the number of experts held
  between forward and backward. Bit-exact by construction — recomputation changes what is saved, never
  what is computed (pinned at rtol=0/atol=0, in and out of grad mode). Cost: one extra dequantize per
  projection in backward; subsumed by full gradient checkpointing when enabled.
• Frozen base: packed weights are requires_grad=False; the forward is differentiable w.r.t. the
  input activations, so it drops in as the frozen base of a QLoRA setup (a per-expert LoRA example lives
  in examples/experts4bit_qlora_demo.py).
• Dtype-cast safety: .to(dtype) / .half() / .bfloat16() retarget compute_dtype; the fp32
  absmax/codebook state is shielded from the cast (a stock nn.Module cast would silently degrade
  every subsequent dequantization — pinned by a test).

Correctness & coverage — tests/test_experts4bit.py, 53 tests, CPU + CUDA (RTX A2000)

• Forward matches a full-precision fused-experts reference built from the same dequantized weights
  (float32, isolating routing from quant error, 1e-4), and is bit-exact vs plain
  dequantize-then-linear (rtol=0/atol=0).
• Backward matches the reference; the frozen base receives no gradient; saved_tensors_hooks
  proves no weight-shaped tensor (either orientation) is saved while a plain dequantize+linear control
  does save one, with gradients matching the control exactly.
• Device movement: .to() carries packed weights, absmax, and codebook together; cpu→cuda→cpu
  round-trips bit-exact (bytes, scales, and forward).
• Serialization: state_dict round-trip bit-exact, strict and non-strict into a ctor-built module;
  safetensors save_model/load_model round-trip with bit-exact forward parity.
• CUDA numerics: per-expert nf4 dequant mean-abs error ≤ 0.008 at the test file's weight scale
  (measures ~0.0073 on an RTX A2000); forward rel err ≤ 0.2 vs the true-weight reference.
• Failed to quant MoE models with fused expert weights in transformers v5 #1849 regression: a fused 3-D nn.Parameter expert module (no nn.Linear for the walker to catch)
  is actually 4-bit-quantized by from_float — uint8-packed, >3× smaller than fp16.
• End-to-end: a frozen Experts4bit base + trainable per-expert LoRA reduces MSE over 30 steps while
  the 4-bit base stays bit-identical.

Non-goals (kept out to keep the diff strictly 4-bit)

• N-bit storage generalization (int8/fp8/bf16/fp16 under the same layout) — implemented and tested on
  the fork's experts-nbit-followup
  branch; happy to open it as a follow-up once this lands.
• Grouped-GEMM — the forward is a per-expert loop; a fused grouped kernel is future work.
• compress_statistics / double-quant — deliberately not used; absmax stays a plain fp32 buffer
  (no nested state). Straightforward to add later.
• CPU-offload / streaming loader — a separate capability that lives out-of-tree (in the
  experts4bit-qlora package); this PR is the quantization primitive only.

Notes

• Pure Python over the stable bitsandbytes.functional APIs
  (quantize_4bit / dequantize_4bit / get_4bit_type / QuantState) — no kernel changes, so the
  tests pass against released bitsandbytes as well as source; nothing here depends on unreleased kernels.
• It's a memory optimization (experts 16-bit → 4-bit), not a throughput one — the per-expert dequant
  loop trades compute for memory, consistent with how bitsandbytes frames 4-bit generally.

Closes #1849. cc @matthewdouglas @SunMarc

[matthewdouglas]

Hi, thanks for the PR. I am a little concerned with how quickly it was opened after discussion. With that said I'll follow up soon, but likely we won't merge something for this until after v0.50.0 release.

[pjordanandrsn]

Thanks @matthewdouglas — fair concern. The asking-first part was real: nothing was written until the shape was pinned down, and the PR follows it — plain nn.Parameter + per-expert absmax buffers on the module, compress_statistics deferred, in_features % blocksize == 0 enforced, per-expert dequant loop for the first cut. The footprint/VRAM numbers in the description are measured on my own A2000 12 GB; fitting fused-experts MoE models on a 12 GB card is the itch that started this.

No rush from me; post-v0.50.0 was always the plan. Converting to draft so it reads as what it is — something concrete to react to when you pick the feature up. Happy to rework it toward whatever you land on, or for you to cherry-pick the useful parts.

[David-AU-github]

So looking forward to this; this will allow Unsloth training of Qwen 3.5/.6 and Gemma 4 sparse moes on reasonable hardware vs H100/200.

[pjordanandrsn]

Updated this branch with backward/gradient + QLoRA-training test coverage and a small OLMoE-1B-7B demo validating quant-correctness and the ~4.7 GB footprint on a 12 GB card (per your June note). Still the correctness-first per-expert loop — no grouped-GEMM.

No rush given v0.50.0 isn't out yet; happy to leave this in draft until you pick it up.

Separately I have a memory-lowering matmul_4bit forward (gated so it only engages on bitsandbytes ≥ 0.50, else falls back to the bit-exact dequantize path) — glad to add it here, keep it as a follow-up, or leave it for your MoE-kernel work, whichever you prefer.

[pjordanandrsn]

Rebased on current main; added a regression test for the exact #1849 silent-skip (fused 3-D module with no nn.Linear → asserts from_float actually 4-bit-quantizes it) plus shape-parametrized forward coverage — 38 tests, CPU+CUDA. Still draft, no rush.

[pjordanandrsn]

@matthewdouglas, a couple updates:

Training path no longer needs v0.50.0: backward is now a small _FrozenLinearRecomputeBackward autograd Function—re-dequantizes each expert instead of saving it, bit-exact gradients on any current bitsandbytes, no matmul_4bit gate. Purely additive: nn/experts.py + export, docs, example, tests. A downstream QLoRA package built on it runs on stock bitsandbytes>=0.43 (on PyPI, in-tree Kaggle-T4 notebook).

Also folded Experts4bit into an ExpertsNbit base (4-bit + 8/16-bit storage)—past this PR's 4-bit title, so I can split the N-bit into a follow-up or leave it, your call.

Still draft, no rush.

[pjordanandrsn]

Marking ready for review — scope is strictly 4-bit (the N-bit generalization moved to the fork's experts-nbit-followup branch), and the training path needs no unreleased kernels: _FrozenLinearRecomputeBackward re-dequantizes per expert in backward, bit-exact on stock bitsandbytes. Since the draft: device-movement, dtype-cast, strict/non-strict state_dict, safetensors, and CUDA-numerics tests — 53 tests green on CPU and an RTX A2000. Built to the plain-nn.Parameter + module-buffer spec from #1849 — happy to adjust the API surface.