OS-Inspired Memory Management for MLX Training on Apple Silicon
mlx-pager brings operating system memory management techniques to MLX model training, enabling training of models that exceed physical memory on Apple Silicon Macs.
Train a 27B model on a 24GB Mac. Train a 70B model on a 64GB Mac. No NVIDIA GPU required.
Apple Silicon Macs have fast unified memory and NVMe SSDs (5-7 GB/s), but MLX has no mechanism to use SSD as memory extension during training. When a model's training footprint exceeds physical RAM, you get an OOM crash.
| Model | Training Memory Needed | 24GB Mac | 64GB Mac |
|---|---|---|---|
| 3B 4bit + LoRA | ~11 GB | Works | Works |
| 9B 4bit + LoRA | ~18 GB | OOM | Works |
| 14B 4bit + LoRA | ~28 GB | OOM | OOM |
| 27B 4bit + LoRA | ~45 GB | OOM | OOM |
| 35B-A3B MoE 4bit | ~25 GB | OOM | OOM |
With mlx-pager:
| Model | Physical Memory | mlx-pager | Status |
|---|---|---|---|
| 9B 4bit + LoRA | 12 GB available | ~8 GB | Works |
| 14B 4bit + LoRA | 12 GB available | ~10 GB | Works |
| 27B 4bit + LoRA | 50 GB available | ~15 GB | Works |
| 70B 4bit + LoRA | 50 GB available | ~25 GB | Works |
mlx-pager treats Apple Silicon's NVMe SSD as a "swap partition" for model training, implementing OS-level memory management concepts:
┌─────────────────────────────────────────┐
│ mlx-pager Memory Hierarchy │
├─────────────────────────────────────────┤
│ │
│ L0: MLX Memory (Hot) │
│ ├── LoRA adapters (always resident) │
│ ├── Optimizer states (always resident) │
│ ├── Embedding + LM Head (pinned) │
│ └── Current layer weights + activations│
│ Bandwidth: ~400 GB/s │
│ │
│ L1: NVMe SSD (Warm) │
│ ├── Frozen layer weights │
│ ├── Checkpointed activations │
│ └── Prefetched next-layer weights │
│ Bandwidth: ~5-7 GB/s │
│ │
│ L2: Disk (Cold) │
│ └── Full model checkpoint │
│ Bandwidth: ~0.5-2 GB/s │
│ │
└─────────────────────────────────────────┘
| OS Concept | mlx-pager Implementation |
|---|---|
| Virtual Memory / MMU | Model "sees" unlimited memory; pager manages physical ↔ SSD mapping |
| Page Tables | Layer registry tracks which layers are in memory vs. on SSD |
| Page Fault Handler | When a layer is needed but not in memory, load from SSD |
| LRU Page Replacement | Evict least-recently-used frozen layers to SSD |
| Prefetch / Readahead | Asynchronously load next layer while current layer computes |
| TLB (Hot Cache) | Pin frequently-accessed layers (Embedding, LM Head, LoRA) |
| mlock | Mark critical layers as non-evictable |
| Copy-on-Write | LoRA/DoRA = CoW for model weights |
| Demand Paging | Gradient checkpointing = recompute instead of storing |
| Swap Priority | Configurable eviction order based on layer access patterns |
| OOM Killer | Graceful degradation: auto-reduce batch/seq/layers before crashing |
| Huge Pages | Batch I/O operations for large contiguous transfers |
pip install mlx-pager
# Drop-in replacement for mlx-lm training
from mlx_pager import PagedTrainer
trainer = PagedTrainer(
model="mlx-community/Qwen3.5-27B-Instruct-4bit",
memory_budget="12GB", # Your available memory
ssd_cache_dir="/tmp/mlx-pager-cache",
)
trainer.lora(
data="./train_data",
lora_rank=16,
learning_rate=1e-5,
max_seq_length=2048,
)See docs/architecture.md for the full technical design.
See docs/implementation-plan.md for the phased development roadmap.
See docs/performance-targets.md for benchmarks and testing strategy.
See CONTRIBUTING.md for development setup and guidelines.
Apache 2.0
- Inspired by Linux kernel memory management (MMU, swap, LRU, prefetch)
- Built on Apple MLX and mlx-lm
- Informed by DeepSpeed ZeRO-Offload
- Motivated by the AI-Love project's need to train large models on consumer hardware