Zero order training cuts memory use for large AI models fast
Probe-Space Preconditioning for Fast and Stable Zero-Order Training
Machine Learning
Summary
Training very large AI models usually requires a lot of computer memory for storing extra information, which can be a problem. The paper looks at a different kind of training called zero-order optimization that needs much less memory but has been slower. The authors improve zero-order training by changing how many trial guesses are made each step and adding a simple extra check that helps the training focus on important directions. Their new method trains big models more quickly and efficiently than before, using ordinary GPUs.
What this means in practice
- •For deep learning engineers: Train very large AI models efficiently on regular GPUs by using zero-order methods with probe-space preconditioning to reduce memory needs and speed up convergence.
- •For machine learning infrastructure teams: Build distributed training pipelines that handle large models with limited GPU resources by integrating efficient zero-order optimizers featuring low memory footprints.
Authors
Francois Chaubard, Mykel J. Kochenderfer, Chris Ré
Abstract
Backpropagation (BP) dominates deep learning but imposes a massive memory tax. For example, training OPT-30B with Adam requires $\approx$ 600GB of GPU memory (assuming batch size 8 and sequence length 2048). Alternatively, zero-order optimization (ZOO) trains in inference-mode (requiring only $\approx$ 60GB for the same model): no stored activations, no gradients, and no optimizer states. However, ZOO convergence has lagged behind BP. In this work, we evaluate two methods to close this gap. First, we show that reallocating training compute budget from many steps to large effective batch sizes with many perturbations (or probes) but fewer steps, allows 1SPSA (Spall, 1992) to outperform zero order methods like MeZO (Malladi et al., 2023) with less training compute. Next, we introduce 1.5-SPSA, adding a single "clean" forward-pass per step to 1SPSA to calculate a cheap diagonal preconditioner in probe-space, which improves convergence rate and convergence by down-weighting high curvature directions. Benchmarking on 6 post-training datasets on both Qwen3 and OPT model families, we show that 1.5-SPSA achieves State-of-the-Art results over previous ZOO solvers with much less optimization steps. For example, we train OPT-13B (for direct comparison to MeZO) and find 1.5-SPSA achieves +3.1% accuracy on SST-2 over both MeZO and BP in only 70 steps vs. MeZO's 100,000 steps. Finally, we combine an 8-bit-packing random generator, triton fused unpack/apply kernels, and distributed parallelism to achieve fast and stable training of models as large as OPT-30B in-place on commodity GPUs (e.g. A100).