Experts trainability: bit-exact recompute-in-backward + ExpertsNbit (8/16-bit storage)

[pjordanandrsn]

Stages the precision-generalization + trainability work into feature/experts-4bit (the bitsandbytes-foundation#1965 branch), as two commits.

1. feat: generalize Experts4bit → ExpertsNbit (8-bit + 16-bit storage)

One module parametrized by quant_type; Experts4bit stays as a fully back-compatible fixed-4-bit subclass (same constructor, from_float, state_dict layout — all pre-existing tests pass unmodified).

quant_type	storage	compression	notes
nf4 / fp4	4-bit packed (2/byte)	~4×	unchanged 4-bit path
int8	blockwise dynamic map, 1 byte/value	2×	via quantize_blockwise; not LLM.int8() vectorwise (that scheme doesn't map onto the fused 3D stack — possible follow-up discussion)
fp8	e4m3 codebook (create_fp8_map(True, 4, 3, 8))	2×	same blockwise path, float codebook
bf16 / fp16	passthrough (no absmax/codebook)	1×	reference baseline / per-layer opt-out; bnb has no 16-bit quantization, so this is deliberately not a quant scheme

Fidelity: int8 blockwise mean-abs reconstruction error is >4× lower than nf4 on the same weights — locked in by test_expertsnbit_fidelity_ordering so a codebook/blocksize regression fails loudly. (This commit also folds in a one-line ruff-format fix to the base branch's bitsandbytes-foundation#1849 regression test; CI's format check fails without it.)

2. mem: re-dequantize frozen expert weights in backward instead of saving them

Each expert projection goes through a small autograd Function (_FrozenLinearRecomputeBackward) whose forward is dequantize + F.linear and whose backward re-dequantizes the frozen weight to compute grad_output @ W. Since the base is frozen, no weight gradient exists, so the dequantized [out, in] weight never enters autograd's saved-tensor storage — activation memory stays independent of how many experts sit between forward and backward.

• Bit-exact by construction. Recomputation changes what is saved, never what is computed — outputs are bit-identical to plain dequantize+linear in and out of grad mode, on every device. (An earlier iteration routed through bnb.matmul_4bit and was rejected: the fused gemm_4bit kernel it dispatches to for small token batches differs from dequantize+linear by accumulation order.) test_experts4bit_forward_is_bit_exact_dequantize_linear pins the equivalence at rtol=0/atol=0.
• Scheme-agnostic. The recompute closure is just _dequantize_expert, so int8/fp8/passthrough get the same training-memory behavior — no per-scheme kernel needed.
• Cost: one extra dequantize per projection in backward (subsumed by full gradient checkpointing when enabled). In a matched A/B on our bench GPU the step-time delta vs eager was below run-to-run noise — see the measurements below.
• test_experts4bit_backward_saves_no_dequantized_weight locks the memory contract via saved_tensors_hooks: nothing weight-shaped (either orientation) is saved by the module, while a plain dequantize+linear control does save it — and gradients match that control exactly.

Measurements (RTX A2000 12GB)

Matched A/B at OLMoE-1B-7B expert dims — 64 experts, hidden 2048, intermediate 1024, top-8, bf16, 512 tokens, a stack of L frozen Experts4bit layers with gradients flowing to the input activations (the QLoRA regime). The only difference between arms is _project: eager dequantize+linear (weight saved by autograd) vs the recompute Function.

	eager, L=2	Function, L=2	eager, L=4	Function, L=4
peak CUDA memory / step	2121 MB	585 MB	4201 MB	1129 MB
step overhead above resting	1670 MB	134 MB (−92%)	3318 MB	246 MB (−93%)

The eager arm's overhead is exactly the saved dequantized weights (~825 MB/layer at these dims) and grows linearly with depth; the Function arm's ~56 MB/layer is activations only. With ~6 GB free on the card, the eager arm OOMs at L=4 while the Function arm runs. Forward outputs were bit-identical between arms in every run. Step-time differences between the two mechanisms were within run-to-run noise on the (contended) bench host — repeated eager baselines of the same config varied 231–391 ms while the Function arm measured 389–470 ms, with the inference-forward comparison flipping sign between runs — so no reliable wall-clock penalty was resolvable; the memory numbers are allocator-exact and reproduced identically across runs.

Verification

• tests/test_experts4bit.py + tests/test_experts_nbit.py: CPU 42/42 passed on Linux/Python 3.11 and Windows 11/Python 3.13 (both CPU-only torch, no native build); CUDA (RTX A2000, sm_86, self-built dev .so): 34/34 -cuda variants pass, including the Failed to quant MoE models with fused expert weights in transformers v5 bitsandbytes-foundation/bitsandbytes#1849 regression/shape tests from this branch's base and the new bit-exactness/saved-tensors tests.
• pre-commit run --all-files: all hooks pass.
• Frozen-base QLoRA training verified end-to-end on the 4-bit path in earlier work on this branch (OLMoE-1B-7B, held-out eval 1.48→1.03); the int8 path has an equivalent frozen-base LoRA training test.

🤖 Generated with Claude Code

[claude]

⚠️ Code review skipped — your organization has reached its monthly code review spending cap.

An organization admin can view or raise the cap at claude.ai/admin-settings/claude-code. The cap resets at the start of the next billing period.

Once the cap resets or is raised, push a new commit or reopen this pull request to trigger a review.