mem: re-dequantize frozen expert weights in backward instead of saving 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 — and the
  mechanism is scheme-agnostic (nf4/fp4/int8/fp8/passthrough), no
  per-scheme 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>

Author: pjordanandrsn

bitsandbytes/nn/experts.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 from collections.abc import Callable
+import functools
 from typing import Optional
 
 import torch
@@ -42,6 +43,34 @@ def _build_code(quant_type: str, device) -> Optional[torch.Tensor]:
     return code.to(device) if device is not None else code
 
 
+class _FrozenLinearRecomputeBackward(torch.autograd.Function):
+    """``F.linear`` against a frozen dequantized weight, re-dequantizing it in backward.
+
+    The weight produced by ``dequant_fn`` (a closure over the packed buffers) is an
+    intermediate, not a Parameter, so a plain ``F.linear`` would stash it as a saved
+    activation for the whole forward-to-backward window — one full-precision expert
+    weight per projection per layer. Because the base is frozen, backward needs no
+    gradient for the weight and only computes ``grad_output @ weight``; the weight can
+    therefore be dropped after the forward matmul and re-dequantized on demand, keeping
+    training memory independent of the number of experts held between forward and
+    backward. Numerically identical to dequantize-then-``linear`` by construction — the
+    forward *is* dequantize-then-linear; recomputation only changes what is saved, never
+    what is computed.
+    """
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, dequant_fn: Callable[[], torch.Tensor]) -> torch.Tensor:
+        ctx.dequant_fn = dequant_fn
+        return F_nn.linear(x, dequant_fn())
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        grad_x = None
+        if ctx.needs_input_grad[0]:
+            grad_x = grad_output @ ctx.dequant_fn()
+        return grad_x, None
+
+
 class ExpertsNbit(nn.Module):
     """Low-bit quantized storage for fused Mixture-of-Experts expert weights.
 
@@ -72,8 +101,11 @@ class ExpertsNbit(nn.Module):
     ``state_dict`` mechanism with no custom save/load hooks.
 
     The forward pass dequantizes a single expert at a time (a per-expert loop), mirroring
-    the reference fused-experts forward. Grouped-GEMM is intentionally left for future
-    work.
+    the reference fused-experts forward. In training, the dequantized weight is not kept
+    as a saved activation: it is re-dequantized on demand in backward (see
+    :class:`_FrozenLinearRecomputeBackward`), for every storage scheme, so activation
+    memory stays independent of the number of experts. Grouped-GEMM is intentionally left
+    for future work.
 
     <Tip warning={true}>This feature is experimental and may change in future releases.</Tip>
 
@@ -299,9 +331,13 @@ def _dequantize_expert(
         return F.dequantize_4bit(packed[expert_idx].reshape(-1, 1), quant_state=quant_state)
 
     def _project(self, packed, absmax, shape, expert_idx, x, compute_dtype):
-        """One expert projection: dequantize + ``linear``."""
-        weight = self._dequantize_expert(packed, absmax, shape, expert_idx, compute_dtype)
-        return F_nn.linear(x, weight)
+        """One expert projection: dequantize + ``linear``, re-dequantizing in backward.
+
+        Works identically for every storage scheme — the recompute closure is just
+        :meth:`_dequantize_expert` — and never produces a gradient for the frozen storage.
+        """
+        dequant_fn = functools.partial(self._dequantize_expert, packed, absmax, shape, expert_idx, compute_dtype)
+        return _FrozenLinearRecomputeBackward.apply(x, dequant_fn)
 
     def forward(
         self,

tests/test_experts4bit.py

@@ -348,3 +348,90 @@ def test_experts4bit_shapes(device, num_experts, hidden, inter):
     top_k_weights = torch.softmax(torch.randn(n_tok, TOP_K, device=device), dim=-1)
     out = module(hidden_states, top_k_index, top_k_weights)
     assert out.shape == (n_tok, hidden) and torch.isfinite(out).all()
+
+
+# --- recompute-in-backward projection path ----------------------------------------------
+# _project routes every expert projection through _FrozenLinearRecomputeBackward: the forward IS
+# dequantize + F.linear (bit-exact by construction, every device and grad mode), and the backward
+# re-dequantizes the frozen weight on demand instead of keeping it as a saved activation. These
+# tests pin both halves of that contract — numerics (grad mode changes nothing; matches a plain
+# dequantize+linear reference bit-for-bit) and memory (no [out, in] weight is ever saved).
+
+
+def _dequantized_expert_stacks(module):
+    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)
+        ]
+    )
+    return gate_up_deq, down_deq
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+def test_experts4bit_forward_is_bit_exact_dequantize_linear(device):
+    """Forward equals a plain dequantize+linear reference bit-for-bit, in and out of grad mode."""
+    gate_up, down = _random_expert_weights(torch.float32, device)
+    module = Experts4bit.from_float(gate_up, down, compute_dtype=torch.float32)
+    hidden_states, top_k_index, top_k_weights = _random_routing(device)
+
+    gate_up_deq, down_deq = _dequantized_expert_stacks(module)
+    ref = _reference_forward(gate_up_deq, down_deq, hidden_states, top_k_index, top_k_weights)
+
+    out_grad_mode = module(hidden_states, top_k_index, top_k_weights)
+    with torch.no_grad():
+        out_no_grad = module(hidden_states, top_k_index, top_k_weights)
+
+    torch.testing.assert_close(out_grad_mode, ref, rtol=0, atol=0)
+    torch.testing.assert_close(out_no_grad, ref, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+def test_experts4bit_backward_saves_no_dequantized_weight(device):
+    """The dequantized expert weights are dropped after each forward matmul (re-dequantized in
+    backward), so nothing weight-shaped reaches autograd's saved-tensor storage — while a plain
+    dequantize+linear control does save them. Gradients still match the control exactly."""
+    gate_up, down = _random_expert_weights(torch.float32, device)
+    module = Experts4bit.from_float(gate_up, down, compute_dtype=torch.float32)
+    hidden_states, top_k_index, top_k_weights = _random_routing(device)
+    # F.linear's backward may save the weight either as-is or pre-transposed (device/impl
+    # dependent), so match both orientations.
+    weight_shapes = set()
+    for shape in (module._gate_up_shape, module._down_shape):
+        weight_shapes.add(tuple(shape))
+        weight_shapes.add(tuple(reversed(shape)))
+
+    def run_recording_saved_shapes(fn, x):
+        saved = []
+
+        def pack(t):
+            saved.append(tuple(t.shape))
+            return t
+
+        with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
+            out = fn(x)
+        return out, saved
+
+    x_mod = hidden_states.detach().clone().requires_grad_(True)
+    out_mod, saved_mod = run_recording_saved_shapes(lambda x: module(x, top_k_index, top_k_weights), x_mod)
+    assert not (set(saved_mod) & weight_shapes)
+
+    gate_up_deq, down_deq = _dequantized_expert_stacks(module)
+    x_ref = hidden_states.detach().clone().requires_grad_(True)
+    out_ref, saved_ref = run_recording_saved_shapes(
+        lambda x: _reference_forward(gate_up_deq, down_deq, x, top_k_index, top_k_weights), x_ref
+    )
+    assert set(saved_ref) & weight_shapes
+
+    out_mod.sum().backward()
+    out_ref.sum().backward()
+    torch.testing.assert_close(out_mod, out_ref, rtol=0, atol=0)
+    torch.testing.assert_close(x_mod.grad, x_ref.grad, rtol=0, atol=0)