feat: Add prepare_model_for_kbit_training and CPU offload checkpointing

Infrastructure for QLoRA training:

prepare_model_for_kbit_training():
- Freezes all base parameters
- Casts normalization layers to float32
- Enables gradient checkpointing
- Pre-allocates global weight buffer for largest layer

checkpoint_cpu_offload():
- Gradient checkpoint that offloads activations to CPU
- Async non-blocking transfers overlap with GPU compute
- RNG state preservation for dropout reproducibility
- Verified to reduce GPU peak memory vs standard forward

8 new tests (3 for prepare_model, 5 for CPU offload checkpoint).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/nn/__init__.py

@@ -20,6 +20,7 @@
     StableEmbedding,
     SwitchBackLinearBnb,
     _GlobalWeightBuffer,
+    prepare_model_for_kbit_training,
 )
 from .triton_based_modules import (
     StandardLinear,

bitsandbytes/nn/modules.py

@@ -921,6 +921,68 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return out.to(inp_dtype)
 
 
+def prepare_model_for_kbit_training(
+    model: torch.nn.Module,
+    use_gradient_checkpointing: bool = True,
+    gradient_checkpointing_kwargs: Optional[dict] = None,
+) -> torch.nn.Module:
+    """Prepare a model with LinearKbit layers for QLoRA-style training.
+
+    This function:
+    1. Freezes all base model parameters (requires_grad=False)
+    2. Casts LayerNorm and other normalization layers to float32
+    3. Enables gradient checkpointing if requested
+    4. Registers the global weight buffer size from the model's largest layer
+
+    After calling this, add LoRA adapters (or any trainable parameters) and
+    those will be the only parameters that receive gradients.
+
+    Args:
+        model: A model containing LinearKbit layers.
+        use_gradient_checkpointing: Enable gradient checkpointing for memory savings.
+        gradient_checkpointing_kwargs: Kwargs passed to model.gradient_checkpointing_enable().
+
+    Returns:
+        The modified model (in-place).
+    """
+    # Freeze all parameters
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # Cast normalization layers to float32 for training stability
+    for module in model.modules():
+        if isinstance(module, (torch.nn.LayerNorm, torch.nn.RMSNorm)):
+            module.float()
+
+    # Enable gradient checkpointing
+    if use_gradient_checkpointing:
+        if hasattr(model, "gradient_checkpointing_enable"):
+            kwargs = gradient_checkpointing_kwargs or {}
+            model.gradient_checkpointing_enable(**kwargs)
+        elif hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        model.is_gradient_checkpointing = True
+
+    # Register global weight buffer for the largest LinearKbit layer
+    max_elements = 0
+    compute_dtype = torch.float16
+    device = None
+    for module in model.modules():
+        if isinstance(module, LinearKbit) and module.weight.kbit_quantized:
+            w = module.weight
+            n = w.N_padded * w.K_dim
+            if n > max_elements:
+                max_elements = n
+                device = w.packed.device
+            if module.compute_dtype is not None:
+                compute_dtype = module.compute_dtype
+
+    if max_elements > 0 and device is not None:
+        _GlobalWeightBuffer.get_buffer(device, max_elements, compute_dtype)
+
+    return model
+
+
 class Int8Params(torch.nn.Parameter):
     def __new__(
         cls,

bitsandbytes/training.py

@@ -0,0 +1,132 @@
+"""Training utilities for kbit QLoRA.
+
+Provides gradient checkpointing with CPU offload for reducing GPU memory
+during QLoRA fine-tuning.
+"""
+
+from typing import Any
+
+import torch
+
+
+class _CPUOffloadCheckpointFunction(torch.autograd.Function):
+    """Gradient checkpoint that offloads activations to CPU during forward.
+
+    Forward: copies activations to CPU asynchronously, frees GPU copy.
+    Backward: copies activations back from CPU, recomputes the forward pass.
+
+    This saves GPU memory at the cost of CPU→GPU bandwidth during backward.
+    Non-blocking transfers overlap with GPU compute when possible.
+    """
+
+    @staticmethod
+    def forward(ctx, run_function, preserve_rng_state, *args):
+        ctx.run_function = run_function
+        ctx.preserve_rng_state = preserve_rng_state
+
+        # Save RNG state if requested
+        if preserve_rng_state:
+            ctx.fwd_cpu_state = torch.random.get_rng_state()
+            ctx.had_cuda = torch.cuda._initialized
+            if ctx.had_cuda:
+                ctx.fwd_gpu_state = torch.cuda.get_rng_state()
+
+        # Save inputs to CPU (async)
+        ctx.cpu_inputs = []
+        ctx.input_requires_grad = []
+        for arg in args:
+            if isinstance(arg, torch.Tensor):
+                ctx.input_requires_grad.append(arg.requires_grad)
+                # Async copy to CPU, pin memory for faster D2H transfer
+                cpu_tensor = torch.empty(
+                    arg.shape, dtype=arg.dtype, device="cpu", pin_memory=True,
+                )
+                cpu_tensor.copy_(arg, non_blocking=True)
+                ctx.cpu_inputs.append(cpu_tensor)
+            else:
+                ctx.input_requires_grad.append(None)
+                ctx.cpu_inputs.append(arg)
+
+        # Run the function
+        with torch.no_grad():
+            outputs = run_function(*args)
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        # Restore inputs from CPU (async)
+        inputs = []
+        for cpu_input, req_grad in zip(ctx.cpu_inputs, ctx.input_requires_grad):
+            if isinstance(cpu_input, torch.Tensor):
+                # Async copy back to GPU
+                gpu_tensor = cpu_input.to("cuda", non_blocking=True)
+                if req_grad:
+                    gpu_tensor.requires_grad_(True)
+                inputs.append(gpu_tensor)
+            else:
+                inputs.append(cpu_input)
+
+        # Synchronize to ensure transfers are complete
+        torch.cuda.current_stream().synchronize()
+
+        # Restore RNG state and recompute forward
+        if ctx.preserve_rng_state:
+            rng_devices = []
+            if ctx.had_cuda:
+                rng_devices.append("cuda")
+            with torch.random.fork_rng(devices=rng_devices, enabled=ctx.preserve_rng_state):
+                torch.random.set_rng_state(ctx.fwd_cpu_state)
+                if ctx.had_cuda:
+                    torch.cuda.set_rng_state(ctx.fwd_gpu_state)
+                with torch.enable_grad():
+                    outputs = ctx.run_function(*inputs)
+        else:
+            with torch.enable_grad():
+                outputs = ctx.run_function(*inputs)
+
+        if isinstance(outputs, torch.Tensor):
+            outputs = (outputs,)
+
+        # Compute gradients
+        input_grads = torch.autograd.grad(
+            outputs,
+            [inp for inp in inputs if isinstance(inp, torch.Tensor) and inp.requires_grad],
+            grad_outputs=grad_outputs,
+        )
+
+        # Map gradients back to original input positions
+        grad_iter = iter(input_grads)
+        result = [None, None]  # for run_function and preserve_rng_state
+        for cpu_input, req_grad in zip(ctx.cpu_inputs, ctx.input_requires_grad):
+            if isinstance(cpu_input, torch.Tensor) and req_grad:
+                result.append(next(grad_iter))
+            else:
+                result.append(None)
+
+        # Free CPU copies
+        ctx.cpu_inputs = None
+
+        return tuple(result)
+
+
+def checkpoint_cpu_offload(
+    function: Any,
+    *args: Any,
+    preserve_rng_state: bool = True,
+) -> Any:
+    """Gradient checkpoint with CPU offload.
+
+    Like ``torch.utils.checkpoint.checkpoint`` but offloads saved activations
+    to CPU during forward to reduce GPU memory.  Activations are copied back
+    from CPU asynchronously during backward.
+
+    Args:
+        function: The function to checkpoint.
+        *args: Arguments to the function. Tensors will be offloaded.
+        preserve_rng_state: Preserve and restore RNG state during recompute.
+
+    Returns:
+        Output of the function.
+    """
+    return _CPUOffloadCheckpointFunction.apply(function, preserve_rng_state, *args)

tests/test_linear_kbit.py

@@ -16,7 +16,7 @@
 
 import bitsandbytes as bnb
 from bitsandbytes import _ops  # noqa: F401 — ensure ops are registered
-from bitsandbytes.nn import LinearKbit, ParamsKbit, _GlobalWeightBuffer
+from bitsandbytes.nn import LinearKbit, ParamsKbit, _GlobalWeightBuffer, prepare_model_for_kbit_training
 
 # Skip all tests if CUDA not available
 pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@@ -309,3 +309,40 @@ def test_gradient_all_k(self, k):
         loss.backward()
         assert x.grad is not None
         assert x.grad.shape == x.shape
+
+
+class TestPrepareModelForKbitTraining:
+    """Tests for prepare_model_for_kbit_training."""
+
+    def _make_model(self):
+        """Create a simple model with LinearKbit layers."""
+        model = torch.nn.Sequential(
+            LinearKbit(256, 128, bias=True, k=4),
+            torch.nn.LayerNorm(128),
+            LinearKbit(128, 64, bias=True, k=4),
+        ).to("cuda")
+        return model
+
+    def test_freezes_all_params(self):
+        """All parameters should be frozen after prepare."""
+        model = self._make_model()
+        prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
+        for param in model.parameters():
+            assert not param.requires_grad
+
+    def test_layernorm_float32(self):
+        """LayerNorm should be cast to float32."""
+        model = self._make_model()
+        prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
+        ln = model[1]
+        assert ln.weight.dtype == torch.float32
+
+    def test_buffer_pre_allocated(self):
+        """Global weight buffer should be sized for the largest layer."""
+        _GlobalWeightBuffer.clear()
+        model = self._make_model()
+        prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
+        # Largest layer is 256→128: K_dim=256, N_padded=128, so 256*128 = 32768
+        buf = _GlobalWeightBuffer._buffers.get(torch.device("cuda", 0))
+        assert buf is not None
+        assert buf.numel() >= 32768

tests/test_training.py

@@ -0,0 +1,121 @@
+"""
+Tests for training utilities (gradient checkpointing with CPU offload).
+
+Verifies:
+- Correctness: output matches standard forward/backward
+- Memory reduction: GPU memory is lower with CPU offload
+- Gradient flow: gradients propagate correctly through checkpoint
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from bitsandbytes.training import checkpoint_cpu_offload
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _simple_block(x):
+    """A simple compute block for testing."""
+    return torch.nn.functional.gelu(x @ x.t()) @ x
+
+
+class TestCPUOffloadCheckpoint:
+    """Tests for checkpoint_cpu_offload."""
+
+    def test_forward_correctness(self):
+        """Output should match standard (non-checkpointed) forward."""
+        x = torch.randn(4, 64, dtype=torch.float32, device="cuda", requires_grad=True)
+        ref = _simple_block(x.detach().clone().requires_grad_(True))
+        out = checkpoint_cpu_offload(_simple_block, x)
+        diff = (out - ref).abs().max().item()
+        assert diff < 1e-5, f"Forward diff: {diff}"
+
+    def test_gradient_correctness(self):
+        """Gradients should match standard backward."""
+        # Standard
+        x_std = torch.randn(4, 64, dtype=torch.float32, device="cuda", requires_grad=True)
+        out_std = _simple_block(x_std)
+        out_std.sum().backward()
+        grad_std = x_std.grad.clone()
+
+        # Checkpointed
+        x_ckpt = x_std.detach().clone().requires_grad_(True)
+        out_ckpt = checkpoint_cpu_offload(_simple_block, x_ckpt)
+        out_ckpt.sum().backward()
+        grad_ckpt = x_ckpt.grad.clone()
+
+        diff = (grad_std - grad_ckpt).abs().max().item()
+        assert diff < 1e-5, f"Gradient diff: {diff}"
+
+    def test_with_nn_module(self):
+        """Should work with nn.Module as the function."""
+        linear = nn.Linear(64, 64).cuda()
+
+        x = torch.randn(4, 64, dtype=torch.float32, device="cuda", requires_grad=True)
+        out = checkpoint_cpu_offload(linear, x)
+        out.sum().backward()
+        assert x.grad is not None
+        assert x.grad.shape == x.shape
+
+    def test_memory_reduction(self):
+        """CPU offload should use less GPU memory than standard checkpoint."""
+        dim = 1024
+
+        # Standard forward (saves activations on GPU)
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+
+        layers = nn.ModuleList([nn.Linear(dim, dim).cuda() for _ in range(4)])
+        x = torch.randn(32, dim, device="cuda", requires_grad=True)
+
+        # Standard: all activations stay on GPU
+        h = x
+        for layer in layers:
+            h = torch.nn.functional.gelu(layer(h))
+        h.sum().backward()
+        peak_standard = torch.cuda.max_memory_allocated()
+
+        # Reset
+        del h, x
+        for p in layers.parameters():
+            p.grad = None
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+
+        # CPU offload: activations go to CPU
+        x = torch.randn(32, dim, device="cuda", requires_grad=True)
+        h = x
+        for layer in layers:
+            h = checkpoint_cpu_offload(lambda inp, l=layer: torch.nn.functional.gelu(l(inp)), h)
+        h.sum().backward()
+        peak_offload = torch.cuda.max_memory_allocated()
+
+        # CPU offload should use less peak memory
+        # Allow some margin since PyTorch internal allocations vary
+        assert peak_offload < peak_standard, (
+            f"CPU offload ({peak_offload / 1e6:.1f} MB) should use less peak memory "
+            f"than standard ({peak_standard / 1e6:.1f} MB)"
+        )
+
+    def test_preserves_rng_state(self):
+        """RNG state should be preserved for dropout reproducibility."""
+        linear = nn.Linear(64, 64).cuda()
+        dropout = nn.Dropout(0.5)
+
+        def block(x):
+            return dropout(linear(x))
+
+        torch.manual_seed(42)
+        x = torch.randn(4, 64, device="cuda", requires_grad=True)
+
+        # Run twice with same seed — should produce same output
+        torch.manual_seed(123)
+        out1 = checkpoint_cpu_offload(block, x)
+
+        torch.manual_seed(123)
+        out2 = checkpoint_cpu_offload(block, x.detach().clone().requires_grad_(True))
+
+        diff = (out1 - out2).abs().max().item()
+        assert diff < 1e-6, f"RNG state not preserved: diff={diff}"