test: add backward + QLoRA-training coverage and demo for Experts4bit

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>

Author: pjordanandrsn

examples/experts4bit_qlora_demo.py

@@ -0,0 +1,154 @@
+"""QLoRA-style training of fused MoE experts on a frozen ``Experts4bit`` base.
+
+This is a *reference pattern*, intentionally **not** part of the bitsandbytes public API. It
+shows that the ``Experts4bit`` 4-bit storage primitive can serve as the frozen base of a
+QLoRA-style fine-tune of fused Mixture-of-Experts weights: the 4-bit expert weights stay
+frozen, and small per-expert low-rank (LoRA) adapters are the only trainable parameters.
+
+The adapter wiring shown here is the kind of thing that would ultimately live in PEFT /
+Unsloth rather than in bitsandbytes itself — the point of this file is to demonstrate that
+the base primitive is already differentiable and trainable as a frozen base today.
+
+Run:
+    python examples/experts4bit_qlora_demo.py
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from bitsandbytes.nn import Experts4bit
+
+
+class ExpertsLoRA(nn.Module):
+    """Per-expert LoRA adapters over a frozen :class:`Experts4bit` base.
+
+    For each expert ``e``, the two frozen 4-bit projections are augmented with a trainable
+    low-rank term ``scaling * (x @ A[e].T) @ B[e].T``:
+
+      * ``gate_up``: ``A[e]`` is ``[r, hidden]``, ``B[e]`` is ``[gate_up_out, r]``
+      * ``down``:    ``A[e]`` is ``[r, intermediate]``, ``B[e]`` is ``[hidden, r]``
+
+    ``B`` is initialised to zero, so the adapted module is identical to the frozen base at
+    step 0 and only departs from it as the adapters train (standard LoRA initialisation).
+    """
+
+    def __init__(self, base: Experts4bit, r: int = 8, alpha: int = 16, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad_(False)
+
+        self.r = r
+        self.scaling = alpha / r
+
+        num_experts = base.num_experts
+        gate_up_out, hidden = base._gate_up_shape  # [2*intermediate (or intermediate), hidden]
+        _, intermediate = base._down_shape  # [hidden, intermediate]
+
+        self.gate_up_lora_A = nn.Parameter(torch.empty(num_experts, r, hidden, dtype=dtype))
+        self.gate_up_lora_B = nn.Parameter(torch.zeros(num_experts, gate_up_out, r, dtype=dtype))
+        self.down_lora_A = nn.Parameter(torch.empty(num_experts, r, intermediate, dtype=dtype))
+        self.down_lora_B = nn.Parameter(torch.zeros(num_experts, hidden, r, dtype=dtype))
+
+        # A ~ small random, B = 0  =>  the initial LoRA delta is exactly zero.
+        nn.init.normal_(self.gate_up_lora_A, std=1.0 / r)
+        nn.init.normal_(self.down_lora_A, std=1.0 / r)
+
+    def _lora(self, x: torch.Tensor, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+        # x: [n, in]; A: [r, in]; B: [out, r]  ->  [n, out]
+        return self.scaling * F.linear(F.linear(x, A), B)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        base = self.base
+        compute_dtype = base.compute_dtype if base.compute_dtype is not None else hidden_states.dtype
+        hidden_states = hidden_states.to(compute_dtype)
+
+        final_hidden_states = torch.zeros_like(hidden_states, dtype=torch.float32)
+
+        with torch.no_grad():
+            expert_mask = F.one_hot(top_k_index, num_classes=base.num_experts).permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False).view(-1)
+
+        for expert_idx in expert_hit:
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            x = hidden_states[token_idx]
+
+            # Frozen 4-bit base projection + trainable low-rank delta.
+            gate_up_w = base._dequantize_expert(
+                base.gate_up_proj, base.gate_up_absmax, base._gate_up_shape, expert_idx, compute_dtype
+            )
+            proj = F.linear(x, gate_up_w) + self._lora(
+                x, self.gate_up_lora_A[expert_idx], self.gate_up_lora_B[expert_idx]
+            )
+
+            if base.has_gate:
+                gate, up = proj.chunk(2, dim=-1)
+                current_hidden = base.act_fn(gate) * up
+            else:
+                current_hidden = base.act_fn(proj)
+
+            down_w = base._dequantize_expert(
+                base.down_proj, base.down_absmax, base._down_shape, expert_idx, compute_dtype
+            )
+            current_hidden = F.linear(current_hidden, down_w) + self._lora(
+                current_hidden, self.down_lora_A[expert_idx], self.down_lora_B[expert_idx]
+            )
+
+            current_hidden = current_hidden * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden.to(final_hidden_states.dtype))
+
+        return final_hidden_states.to(hidden_states.dtype)
+
+
+def main() -> None:
+    torch.manual_seed(0)
+
+    num_experts, hidden, intermediate = 8, 128, 256
+    num_tokens, top_k = 64, 2
+
+    # A full-precision fused-expert stack (the shape transformers v5 stores MoE experts in).
+    gate_up = torch.randn(num_experts, 2 * intermediate, hidden) * 0.1
+    down = torch.randn(num_experts, hidden, intermediate) * 0.1
+
+    # Freeze it in 4-bit, then attach trainable LoRA adapters.
+    base = Experts4bit.from_float(gate_up, down, quant_type="nf4", compute_dtype=torch.float32)
+    model = ExpertsLoRA(base, r=8, alpha=16)
+
+    trainable = [p for p in model.parameters() if p.requires_grad]
+    n_train = sum(p.numel() for p in trainable)
+    n_base_bytes = base.gate_up_proj.numel() + base.down_proj.numel()
+    print(f"trainable LoRA params: {n_train:,}   frozen packed base bytes: {n_base_bytes:,}")
+
+    hidden_states = torch.randn(num_tokens, hidden)
+    top_k_index = torch.randint(0, num_experts, (num_tokens, top_k))
+    top_k_weights = torch.softmax(torch.randn(num_tokens, top_k), dim=-1)
+    target = torch.randn(num_tokens, hidden)
+
+    gate_up_before = base.gate_up_proj.clone()
+
+    optimizer = torch.optim.Adam(trainable, lr=1e-2)
+    print("\nstep   loss")
+    for step in range(50):
+        optimizer.zero_grad()
+        out = model(hidden_states, top_k_index, top_k_weights)
+        loss = F.mse_loss(out, target)
+        loss.backward()
+        assert base.gate_up_proj.grad is None, "frozen base must never receive a gradient"
+        optimizer.step()
+        if step % 10 == 0 or step == 49:
+            print(f"{step:4d}   {loss.item():.5f}")
+
+    assert torch.equal(base.gate_up_proj, gate_up_before), "frozen base bytes must be unchanged"
+    print("\nbase packed weights unchanged after training:", torch.equal(base.gate_up_proj, gate_up_before))
+
+
+if __name__ == "__main__":
+    main()

tests/test_experts4bit.py

@@ -154,3 +154,141 @@ def test_experts4bit_blocksize_validation():
 
 def test_experts4bit_is_exported():
     assert bnb.nn.Experts4bit is Experts4bit
+
+
+# --- Backward / autograd ---------------------------------------------------------------
+# Experts4bit is a *frozen* 4-bit base: the packed weights are requires_grad=False, so they
+# never receive gradients, but the per-expert dequant + linear + index_add_ forward is fully
+# differentiable w.r.t. the input activations. That makes it usable as the frozen base of a
+# QLoRA-style setup (gradients flow to adapters/earlier layers, not to the quantized weights).
+# These tests lock that contract in.
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype)
+def test_experts4bit_backward_flows_to_input(device, dtype):
+    gate_up, down = _random_expert_weights(dtype, device)
+    module = Experts4bit.from_float(gate_up, down, compute_dtype=dtype)
+
+    hidden_states, top_k_index, top_k_weights = _random_routing(device)
+    hidden_states = hidden_states.to(dtype).detach().requires_grad_(True)
+
+    out = module(hidden_states, top_k_index, top_k_weights)
+    out.float().sum().backward()
+
+    # Gradient reaches the input activations, is finite, and is nonzero (every token is routed
+    # to TOP_K experts here, so every row contributes).
+    assert hidden_states.grad is not None
+    assert torch.isfinite(hidden_states.grad).all()
+    assert hidden_states.grad.float().abs().sum() > 0
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+def test_experts4bit_base_weights_stay_frozen(device):
+    gate_up, down = _random_expert_weights(torch.float32, device)
+    module = Experts4bit.from_float(gate_up, down, compute_dtype=torch.float32)
+
+    # Packed weights are frozen by construction ...
+    assert module.gate_up_proj.requires_grad is False
+    assert module.down_proj.requires_grad is False
+
+    hidden_states, top_k_index, top_k_weights = _random_routing(device)
+    hidden_states = hidden_states.requires_grad_(True)
+    module(hidden_states, top_k_index, top_k_weights).sum().backward()
+
+    # ... and a backward pass leaves no gradient on them (so an optimizer can never nudge the
+    # quantized base, and the absmax buffers are not trainable either).
+    assert module.gate_up_proj.grad is None
+    assert module.down_proj.grad is None
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+def test_experts4bit_backward_matches_reference(device):
+    # float32 throughout: the module's autograd path must match a plain full-precision forward
+    # built from the *same* dequantized weights, isolating gradient correctness from quant error.
+    gate_up, down = _random_expert_weights(torch.float32, device)
+    module = Experts4bit.from_float(gate_up, down, compute_dtype=torch.float32)
+
+    gate_up_deq = torch.stack(
+        [
+            module._dequantize_expert(
+                module.gate_up_proj, module.gate_up_absmax, module._gate_up_shape, e, torch.float32
+            )
+            for e in range(NUM_EXPERTS)
+        ]
+    )
+    down_deq = torch.stack(
+        [
+            module._dequantize_expert(module.down_proj, module.down_absmax, module._down_shape, e, torch.float32)
+            for e in range(NUM_EXPERTS)
+        ]
+    )
+
+    hidden_states, top_k_index, top_k_weights = _random_routing(device)
+    hs_mod = hidden_states.detach().clone().requires_grad_(True)
+    hs_ref = hidden_states.detach().clone().requires_grad_(True)
+
+    out_mod = module(hs_mod, top_k_index, top_k_weights)
+    out_ref = _reference_forward(gate_up_deq, down_deq, hs_ref, top_k_index, top_k_weights)
+
+    out_mod.sum().backward()
+    out_ref.sum().backward()
+
+    torch.testing.assert_close(out_mod, out_ref, rtol=1e-4, atol=1e-4)
+    torch.testing.assert_close(hs_mod.grad, hs_ref.grad, rtol=1e-4, atol=1e-4)
+
+
+def _load_experts_lora():
+    """Load the ExpertsLoRA reference wrapper from examples/ (kept out of the bnb API)."""
+    import importlib.util
+    import os
+
+    path = os.path.join(os.path.dirname(__file__), "..", "examples", "experts4bit_qlora_demo.py")
+    spec = importlib.util.spec_from_file_location("experts4bit_qlora_demo", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.ExpertsLoRA
+
+
+def test_experts4bit_lora_training_reduces_loss():
+    # End-to-end QLoRA-style step: a frozen 4-bit Experts4bit base + trainable per-expert LoRA.
+    # Proves the primitive supports training today — only the adapters move, the base stays put.
+    torch.manual_seed(0)
+    experts_lora = _load_experts_lora()
+
+    gate_up, down = _random_expert_weights(torch.float32, "cpu")
+    base = Experts4bit.from_float(gate_up, down, compute_dtype=torch.float32)
+    model = experts_lora(base, r=4, alpha=8)
+
+    # Only LoRA adapters are trainable; the 4-bit base is frozen.
+    trainable_names = [name for name, p in model.named_parameters() if p.requires_grad]
+    assert trainable_names and all("lora" in name for name in trainable_names)
+
+    gate_up_before = base.gate_up_proj.clone()
+    down_before = base.down_proj.clone()
+
+    hidden_states, top_k_index, top_k_weights = _random_routing("cpu")
+    target = torch.randn_like(hidden_states)
+
+    # Standard LoRA init (B=0) => the adapted forward equals the frozen base forward at step 0.
+    torch.testing.assert_close(
+        model(hidden_states, top_k_index, top_k_weights),
+        base(hidden_states, top_k_index, top_k_weights),
+        rtol=1e-5,
+        atol=1e-5,
+    )
+
+    optimizer = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=1e-2)
+    losses = []
+    for _ in range(30):
+        optimizer.zero_grad()
+        loss = torch.nn.functional.mse_loss(model(hidden_states, top_k_index, top_k_weights), target)
+        loss.backward()
+        assert base.gate_up_proj.grad is None and base.down_proj.grad is None
+        optimizer.step()
+        losses.append(loss.item())
+
+    assert losses[-1] < losses[0]  # training reduces loss
+    # The frozen 4-bit base is bit-identical before and after training.
+    assert torch.equal(base.gate_up_proj, gate_up_before)
+    assert torch.equal(base.down_proj, down_before)