airllm

Decomposes large language models (70B+ parameters) into individual transformer layers that are loaded into GPU memory only when needed during forward passes, then unloaded after computation completes. Uses a layer-by-layer execution strategy where each layer is fetched from disk storage, processed with its input activations, and immediately freed, reducing peak memory footprint from full model size to single-layer size. This architectural approach enables 70B models to run on 4GB VRAM without quantization or distillation.

Overlaps disk I/O operations with GPU computation by prefetching the next transformer layer while the current layer is being processed. Uses a background I/O thread that predicts which layer will be needed next and loads it into a staging buffer during the current layer's forward pass, reducing idle GPU time. Achieves approximately 10% inference speed improvement by hiding disk latency behind computation.

Provides utilities to introspect transformer model architectures and automatically extract layer definitions from model configs. Uses config.json inspection to identify layer count, hidden dimensions, attention heads, and other architectural parameters. Supports dynamic layer extraction for models with non-standard layer structures. Enables programmatic access to layer boundaries and architectural metadata.

Applies optional block-wise quantization to model weights only (not activations) to reduce model disk footprint and loading time, offering 4-bit or 8-bit quantization modes. Unlike traditional quantization that quantizes both weights and activations, this approach preserves activation precision during inference, maintaining model accuracy while achieving up to 3x inference speed improvement through reduced I/O overhead. Quantization is applied during model decomposition and stored per-layer on disk.

Provides a unified AutoModel interface that automatically detects model architecture (Llama, ChatGLM, QWen, Baichuan, Mistral, Mixtral, InternLM) from model config and instantiates the appropriate implementation. Includes platform-specific optimizations: uses MLX framework on macOS for native Apple Silicon acceleration, CUDA on NVIDIA GPUs, and ROCm on AMD GPUs. Abstracts away platform differences through a single Python API.

Decomposes full models into individual transformer layers during first run and persists each layer as a separate disk artifact in a structured directory hierarchy. Uses PyTorch's state_dict serialization to save layer weights, biases, and normalization parameters independently. Subsequent runs load layers on-demand from disk without redecomposition. Supports both full-precision and quantized layer storage with metadata tracking.

Provides architecture-specific implementations for 8+ transformer variants (Llama, ChatGLM, QWen, Baichuan, Mistral, Mixtral, InternLM) while exposing a unified inference interface. Each architecture has custom layer definitions that respect model-specific attention mechanisms, activation functions, and normalization schemes. Unified interface handles tokenization, prompt formatting, and output parsing consistently across all supported models.

Provides explicit support for models with extended context windows (e.g., 32K, 100K token contexts) through optimized attention computation and memory management. Handles long sequences by managing KV-cache memory more efficiently during layer-wise inference, avoiding full KV-cache materialization. Supports position interpolation and other long-context techniques at the layer level.

Provides Direct Preference Optimization training framework as an alternative to traditional RLHF with PPO. DPO eliminates the need for a separate reward model by directly optimizing model weights based on preference pairs (chosen vs rejected completions). Implements preference loss computation, gradient accumulation, and training loops optimized for limited GPU memory. Includes dataset preparation utilities for converting preference data into DPO format.

Provides native macOS support through integration with Apple's MLX framework, enabling optimized inference on Apple Silicon (M1/M2/M3) GPUs. Automatically detects macOS platform and routes inference through MLX backend instead of CUDA/ROCm, leveraging Metal Performance Shaders for GPU acceleration. Maintains layer-sharding architecture while using MLX's memory-efficient tensor operations.

Provides a Python inference API that supports both streaming and non-streaming text generation modes. Implements token-by-token generation with configurable sampling strategies (temperature, top-k, top-p), stopping criteria, and output formatting. Handles prompt tokenization, special token insertion, and response parsing automatically. Supports both single-sequence and batch inference patterns through a unified generate() interface.