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 — 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>

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
@@ -13,6 +14,34 @@
 from bitsandbytes.functional import QuantState
 
 
+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 Experts4bit(nn.Module):
     """4-bit quantized storage for fused Mixture-of-Experts expert weights.
 
@@ -33,8 +62,10 @@ class Experts4bit(nn.Module):
     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`), 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>
 
@@ -217,6 +248,16 @@ def _dequantize_expert(
         # transpose back-compat shim — keyed on A.shape[0] == 1 — from firing).
         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``, re-dequantizing in backward.
+
+        The recompute closure is just :meth:`_dequantize_expert`; no gradient is ever
+        produced for the frozen packed storage.
+        """
+        dequant_fn = functools.partial(self._dequantize_expert, packed, absmax, shape, expert_idx, compute_dtype)
+        return _FrozenLinearRecomputeBackward.apply(x, dequant_fn)
+
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -238,20 +279,18 @@ def forward(
             top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
             current_state = hidden_states[token_idx]
 
-            gate_up_w = self._dequantize_expert(
-                self.gate_up_proj, self.gate_up_absmax, self._gate_up_shape, expert_idx, compute_dtype
+            proj = self._project(
+                self.gate_up_proj, self.gate_up_absmax, self._gate_up_shape, expert_idx, current_state, compute_dtype
             )
-            proj = F_nn.linear(current_state, gate_up_w)
             if self.has_gate:
                 gate, up = proj.chunk(2, dim=-1)
                 current_hidden = self.act_fn(gate) * up
             else:
                 current_hidden = self.act_fn(proj)
 
-            down_w = self._dequantize_expert(
-                self.down_proj, self.down_absmax, self._down_shape, expert_idx, compute_dtype
+            current_hidden = self._project(
+                self.down_proj, self.down_absmax, self._down_shape, expert_idx, current_hidden, compute_dtype
             )
-            current_hidden = F_nn.linear(current_hidden, down_w)
             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))
 

tests/test_experts4bit.py

@@ -347,3 +347,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)