bitsandbytes

Implements block-wise quantization (blocksize=256) of optimizer states in Adam8bit, AdamW8bit, and PagedAdamW classes, reducing optimizer memory footprint by ~75% while maintaining training convergence. Uses a five-layer architecture where Layer 1 exposes PyTorch-compatible optim.Optimizer interfaces, Layer 2 manages custom autograd functions for backward passes, Layer 3 implements core quantization algorithms with QuantState management, and Layers 4-5 dispatch to backend-specific CUDA/CPU kernels. Block-wise quantization divides optimizer states into fixed-size blocks, quantizes each block independently with per-block scaling factors, and dequantizes on-the-fly during parameter updates.

Provides 8-bit inference for large language models through Linear8bitLt module that applies vector-wise quantization to weight matrices while preserving high-precision outliers in a separate buffer. Implements a two-tier quantization strategy: most weights are quantized to 8-bit with per-column scaling factors, while outlier columns (detected via threshold-based heuristics) remain in full precision. During forward pass, quantized weights are dequantized on-the-fly, outlier weights are added back, and the computation proceeds in mixed precision (int8 + fp32 for outliers). This achieves ~50% memory reduction for model weights while maintaining inference quality comparable to full-precision models.

Implements efficient matrix multiplication (GEMM) kernels that operate on quantized weights (int8 or int4) while maintaining full-precision activations and outputs. Kernels dequantize weights on-the-fly during computation, perform multiplication in float32, and produce float32 outputs. Supports mixed-precision: weights are int8/int4, activations are float16/float32, and outputs are float32. Optimized CUDA kernels use tensor cores (on modern GPUs) for efficient int8 computation, achieving 2-4x speedup compared to naive dequantize-then-multiply approach. Handles edge cases: non-standard matrix shapes, batch sizes, and quantization block sizes. Integrates with PyTorch's autograd for backward pass.

Integrates with PyTorch's gradient checkpointing (torch.utils.checkpoint) to reduce training memory footprint by trading computation for memory. Gradient checkpointing discards intermediate activations during forward pass and recomputes them during backward pass, reducing peak memory usage by ~30-40%. Works seamlessly with bitsandbytes quantized layers: forward pass uses quantized weights, backward pass recomputes forward pass to get activations, then computes gradients. Enables combining gradient checkpointing with 8-bit optimizers and 4-bit quantization for maximum memory efficiency: 8-bit optimizer saves 75%, 4-bit quantization saves 75%, gradient checkpointing saves 30-40%, totaling ~95% memory reduction.

Provides CPU-optimized implementations of quantization and dequantization operations using SIMD instructions (AVX2, AVX-512) for inference on CPU-only systems. Implements block-wise dequantization with vectorized operations, reducing CPU inference latency by 5-10x compared to naive scalar implementations. Supports int8 and int4 dequantization with per-block scaling factors. CPU kernels are slower than GPU kernels (10-50x slower than CUDA), but enable inference on systems without GPUs (servers, edge devices, laptops). Automatically selected when GPU is unavailable or explicitly requested.

Implements 4-bit quantization of model weights using NF4 (Normal Float 4-bit, information-theoretically optimal for normally distributed weights) or FP4 (standard floating-point 4-bit) data types, combined with LoRA (Low-Rank Adaptation) adapters for parameter-efficient fine-tuning. Uses double quantization to further compress scaling factors, reducing model memory by ~75%. Linear4bit, LinearNF4, and LinearFP4 modules replace standard nn.Linear layers; during forward pass, 4-bit weights are dequantized to float16/float32, multiplied with inputs, and LoRA adapters (low-rank matrices) are added to the output. Backward pass computes gradients only for LoRA parameters and optimizer states, keeping base model frozen. This enables fine-tuning of 70B models on 24GB GPUs.

Implements Layer 4 of the five-layer architecture: dynamic runtime detection and loading of platform-specific compiled binaries (CUDA, CPU, ROCm, Intel XPU) without requiring users to specify backends explicitly. Uses ctypes-based FFI to load .so/.dll files matching the detected CUDA version and GPU architecture; falls back to CPU implementations if GPU libraries unavailable. Operator registration system maps Python function calls (e.g., quantize_blockwise) to corresponding C/CUDA kernel implementations via a registry. This abstraction allows the same Python API to run on NVIDIA GPUs, AMD GPUs, Intel Arc, and CPU without code changes, and enables graceful degradation when hardware-specific optimizations unavailable.

Implements Layer 3 core data structure for managing quantized tensor metadata: QuantState class encapsulates quantized weights, scaling factors (absmax per block/column), data type (NF4/FP4/INT8), and shape information. Provides serialization/deserialization for saving quantized models to disk and loading them back without recomputation. QuantState tracks which tensors are quantized, their quantization parameters, and enables efficient dequantization on-demand. Integrates with PyTorch's state_dict() mechanism for checkpoint saving, allowing quantized models to be saved and loaded like standard PyTorch models. This abstraction decouples quantization logic from neural network modules and enables composable quantization strategies.

Implements Layer 2 custom PyTorch autograd functions (torch.autograd.Function subclasses) that define forward and backward passes for quantized operations. For example, quantized linear layers use custom autograd to compute forward pass with quantized weights (dequantized on-the-fly) and backward pass that computes gradients with respect to full-precision weights, not quantized weights. This enables training with quantized weights while maintaining gradient flow and convergence properties. Autograd functions handle mixed-precision computation: forward pass may use int8/int4 weights, but backward pass uses float32 gradients. Integrates with PyTorch's autograd graph for compatibility with standard training loops, gradient accumulation, and distributed training.

Provides GlobalOptimManager class that coordinates 8-bit optimizer state quantization across distributed training with FSDP (PyTorch's fully sharded data parallel). FSDP shards model parameters and gradients across GPUs; GlobalOptimManager ensures optimizer states are also sharded and quantized consistently. Handles synchronization of quantization metadata (scaling factors, block information) across devices, manages paging of optimizer states to CPU when GPU memory exhausted, and coordinates gradient accumulation across shards. Integrates with FSDP's backward hook system to trigger optimizer updates at the right time without deadlocks or synchronization issues.

Implements NF4 data type specifically designed for quantizing neural network weights that follow approximately normal distributions. NF4 uses 4 bits to represent 16 quantization levels optimized for Gaussian distributions (derived from inverse normal CDF), achieving information-theoretic optimality for normally distributed data. Unlike standard FP4 (which uses uniform floating-point spacing), NF4 allocates more quantization levels near zero and fewer at extremes, matching the distribution of typical neural network weights. Quantization process: compute per-column or per-block statistics (mean, std), map weights to quantization levels using lookup tables, and store only the quantized values and scaling factors. Dequantization reverses the process on-the-fly during inference or training.

Implements secondary quantization of per-block or per-column scaling factors (absmax values) to further reduce model size. In standard quantization, weights are quantized to 4-bit and scaling factors stored in float32 (4 bytes per factor). Double quantization quantizes these scaling factors themselves to 8-bit, reducing their memory footprint by 75%. Process: compute scaling factors for weights (e.g., absmax per 64-weight block), then quantize these scaling factors to 8-bit with their own meta-scaling factors. During dequantization, scaling factors are dequantized first, then used to dequantize weights. This adds one extra dequantization step but reduces total model size by additional 5-10% with minimal accuracy impact.

Implements PagedAdamW optimizer that pages optimizer states between GPU and CPU memory to enable training models larger than GPU VRAM. Maintains a small working set of optimizer states on GPU (for current batch), pages out states to CPU RAM when not needed, and pages them back in when required. Uses asynchronous PCIe transfers to overlap computation with data movement. Tracks which optimizer states are currently on GPU vs CPU, manages a page table, and coordinates with the training loop to ensure correct state is available during parameter updates. This enables training 70B models on 24GB GPUs by using 200GB+ CPU RAM as extended memory.