Reproduce “OASIS: Online Activation Subspace Learning for Memory-Efficient Training” on OpenResearch

About this paper

Sakshi Choudhary, Utkarsh Saxena, Kaushik Roy

OASIS reduces training memory by learning a changing low-dimensional representation of saved activations and reusing it to compress gradients and optimizer states. It updates the representation during training while leaving the forward pass unchanged. On Llama-2 7B fine-tuned for GSM8K, rank-128 OASIS reports 39.10 ± 1.27 accuracy and 49.60 GB peak memory, versus Adam’s 39.37 ± 0.39 and 95.18 GB; this result is specific to that setting.

Key insights

• The low-rank structure of activations can be leveraged to jointly compress activations, gradients, and optimizer states due to the inherent relationship between activation span and weight gradients.
• The optimal activation subspace is not static but continuously evolves during training, necessitating an online and adaptive tracking mechanism rather than fixed or periodically updated bases.
• Maintaining optimizer states within an evolving low-dimensional subspace requires a specific 'projection-aware' transport mechanism to ensure consistent updates and stable optimization.

Summary by alphaXiv.