Add Experts4bit for 4-bit quantization of fused MoE experts - #1965
Add Experts4bit for 4-bit quantization of fused MoE experts#1965pjordanandrsn wants to merge 9 commits into
Conversation
Sorry, something went wrong.
|
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. |
Sorry, something went wrong.
|
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. |
Sorry, something went wrong.
|
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. |
Sorry, something went wrong.
|
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. |
Sorry, something went wrong.
|
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. |
Sorry, something went wrong.
|
@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. |
Sorry, something went wrong.
|
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. |
Sorry, something went wrong.
…ytes-foundation#1849) transformers v5 stores fused MoE experts as a single 3D nn.Parameter (e.g. OlmoeExperts, Qwen3MoeExperts), which the nn.Linear-based 4-bit walker skips. The experts stay in full precision and load_in_4bit barely shrinks the model (issue bitsandbytes-foundation#1849). Experts4bit holds gate_up_proj and down_proj packed in NF4/FP4 as plain nn.Parameter buffers, with per-expert absmax kept on the module itself. The forward pass dequantizes one expert at a time (a per-expert loop), mirroring the reference fused-experts forward. There is no Params4bit tensor-subclass machinery, so the module serializes through the default state_dict with no custom hooks. - from_float() quantizes existing bf16/fp16 expert stacks - enforces in_features % blocksize == 0 for clean per-expert blocking - double-quant (compress_statistics) and grouped-GEMM intentionally deferred for a first cut - tests: quant round-trip, forward vs. full-precision reference, state_dict round-trip, and validation guards
Prove Experts4bit works as a frozen 4-bit QLoRA base. New tests in tests/test_experts4bit.py cover the autograd contract: gradients reach the input activations, the frozen packed weights never receive a gradient, and the backward matches a full-precision reference forward. Add examples/experts4bit_qlora_demo.py with a small per-expert ExpertsLoRA wrapper over a frozen base, plus a test that a real optimizer step reduces loss while the 4-bit base stays bit-identical. The wrapper is a reference pattern (PEFT/Unsloth territory), intentionally not part of the bitsandbytes public API. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…rametrized forward coverage - test_experts4bit_1849_regression_...: build a fused-3D nn.Parameter expert module (the transformers v5 layout the Linear4bit walker skips), assert Experts4bit.from_float actually 4-bit-quantizes it (uint8-packed, >3x smaller than fp16). Guards the exact silent-skip bitsandbytes-foundation#1849 reports. - test_experts4bit_shapes: forward correctness across a spread of (num_experts, hidden, intermediate). Co-Authored-By: Jordan Anderson <paul.jordan.anderson@gmail.com>
…g them Route every expert projection through _FrozenLinearRecomputeBackward, a small autograd Function whose forward computes dequantize + F.linear directly and whose backward re-dequantizes the frozen weight to form grad_output @ W: - Bit-exact by construction, in and out of grad mode, on every device: recomputation changes what is saved, never what is computed. (An earlier iteration routed through bnb.matmul_4bit instead; it was dropped because the fused gemm_4bit kernel it dispatches to for small token batches differs from dequantize+linear by accumulation order.) - The dequantized [out, in] expert weight never enters autograd's saved-tensor storage, so training activation memory stays independent of the number of experts held between forward and backward — no custom kernel needed. - Cost: one extra dequantize per projection in backward; subsumed by full gradient checkpointing when that is enabled. Tests: test_experts4bit_forward_is_bit_exact_dequantize_linear pins forward == plain dequantize+linear at rtol=0/atol=0 (grad and no_grad); test_experts4bit_backward_saves_no_dequantized_weight uses saved_tensors_hooks to assert nothing weight-shaped (either orientation) is saved while a plain dequantize+linear control does save it, and that gradients match the control exactly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…uant state nn.Module dtype casts (.to(dtype), .half(), .bfloat16()) convert every floating-point tensor. The packed uint8 weights are naturally immune, but the fp32 absmax/code buffers would be silently cast, degrading every subsequent dequantization. Override _apply so a float dtype cast retargets compute_dtype while the quantization state stays fp32 (re-derived from the pre-cast tensors — a cast round-trip would be lossy). Device movement is unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…DA fidelity - to()/round-trip movement: all quant-state tensors move together; cpu->cuda->cpu is bit-exact for packed bytes, absmax, and the forward (dequant inputs are integers + fp32 scales, so movement can never perturb the math). - dtype casts (.to(dtype)/.half()/.bfloat16()): compute_dtype retargets, absmax/code stay fp32, dequantized weights bit-identical across the cast. - state_dict strict=False into a ctor-built module: nothing missing, nothing skipped, bit-exact forward parity. - safetensors save_model/load_model round-trip: works with no _extra_state or custom hooks; bit-exact forward parity. - CUDA numerics: nf4 per-expert dequant mean-abs error <= 0.008 for the 0.1-scaled weights used across this file (measures ~0.0073 on an RTX A2000) and forward rel err <= 0.2 vs the true-weight reference. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The four probes a reviewer aims at a custom autograd Function, pinned as
assertions rather than prose:
- torch.compile: Dynamo graph-breaks cleanly on the data-dependent expert
routing (aten.nonzero) and the Function runs eagerly inside the compiled
wrapper — forward and input-grad stay bitwise-equal to eager.
- gradient checkpointing: composes with recompute-in-backward at 3x
dequants per fwd+bwd vs 2x uncheckpointed (+50%, not a multiply),
bit-exact either way; the count is asserted, not documented.
- autocast: the dequant path is untouched (state bit-identical, weights
still materialize in compute_dtype), output dtype does not drift, values
match the non-autocast forward at bf16 precision.
- meta-device: init-under-torch.device('meta') + load_state_dict(assign=
True) — the HF init_empty_weights-style path — materializes packed
weights and absmax with no meta leftovers, frozen flag intact, forward
bit-identical to the source module. Works with zero custom hooks because
the codebook is rebuilt at init rather than serialized.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
v0.50.0 is out, so this is unparked — rebased onto current main. CI has never run here; the workflows are still at action_required. Locally the test file is green on macOS/ARM64, on linux-x64 at torch 2.4.1 and 2.13.0, and on an A2000, plus lint and docs. One flag I'd rather raise myself: nn.parametrize.replace_parameter_4bit already covers much of this (#1849; #1999 fixed its checkpointing case). Experts4bit adds per-expert dequant instead of the whole fused stack, and an explicit recompute-in-backward contract. If you'd rather have those in parametrize, I'll close this and do that instead. No rush. |
Sorry, something went wrong.
|
@matthewdouglas this implements the approach from the #1849 thread — packed weights as plain nn.Parameters with per-expert absmax as module buffers (no Params4bit subclassing), in_features % blocksize enforced at construction, no compress_statistics. Purely additive, +1218/−0. The four workflows have been at action_required since the Jul 26 rebase — mind approving them so there's real CI to review against? |
Sorry, something went wrong.
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.
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).
be a multiple of blocksize so blocks never straddle an expert boundary (validated).
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.
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.
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).
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)
(float32, isolating routing from quant error, 1e-4), and is bit-exact vs plain
dequantize-then-linear (rtol=0/atol=0).
proves no weight-shaped tensor (either orientation) is saved while a plain dequantize+linear control
does save one, with gradients matching the control exactly.
round-trips bit-exact (bytes, scales, and forward).
safetensors save_model/load_model round-trip with bit-exact forward parity.
(measures ~0.0073 on an RTX A2000); forward rel err ≤ 0.2 vs the true-weight reference.
is actually 4-bit-quantized by from_float — uint8-packed, >3× smaller than fp16.
the 4-bit base stays bit-identical.
Non-goals (kept out to keep the diff strictly 4-bit)
the fork's experts-nbit-followup
branch; happy to open it as a follow-up once this lands.
(no nested state). Straightforward to add later.
experts4bit-qlora package); this PR is the quantization primitive only.
Notes
(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.
loop trades compute for memory, consistent with how bitsandbytes frames 4-bit generally.
Closes #1849. cc @matthewdouglas @SunMarc