BAS: Bridging Adam and SignSGD for Memory-Efficient LLM Training
Abstract
We propose Block Adaptive Signum (BAS), which bridges Adam and SignSGD via block-wise scaling of sign updates. By discarding element-wise second moments, BAS reduces memory overhead relative to AdamW without sacrificing performance in our tested settings. Crucially, BAS mimics Adam’s dynamics closely enough to directly inherit its hyperparameters, matching the performance of AdamW without the need for re‑tuning, a common fragility of prior low‑memory optimizers. This structural alignment makes it particularly suitable for tuning Adam-pretrained models. Furthermore, we exploit the inherent robustness of sign-based updates to store the first moment in FP8 without performance degradation. This shrinks the optimizer‑state footprint to 12.5\% of AdamW’s. We theoretically prove convergence under standard assumptions and introduce a communication-efficient variant enabled by the sign-based update. Across extensive evaluations, including pre‑training a 1.5B model on 100B tokens and supervised fine-tuning of models up to 32B parameters, we demonstrate that BAS achieves performance on par with AdamW.
Lay Summary
Training large AI language models demands enormous computer memory, and a surprising amount of it is not used by the model itself but by the optimizer — the algorithm that decides how to adjust billions of settings during learning. This memory cost forces researchers to use expensive hardware and limits who can participate in cutting-edge AI development. We introduce Block Adaptive Signum (BAS), an optimizer that reduces this memory overhead to just one-eighth of the standard approach. Instead of tracking a separate adjustment statistic for every individual model parameter, BAS groups parameters into blocks and shares a single statistic per group. This also allows the remaining stored values to be compressed into an extremely compact format without hurting performance. Importantly, BAS works as a drop-in replacement for the widely used AdamW optimizer, meaning researchers can switch to it without re-tuning their training setup — a costly process that has limited the adoption of prior memory-efficient alternatives. We validated BAS on models with up to 32 billion parameters across math, coding, and reasoning benchmarks, matching the standard optimizer's performance throughout. By lowering the hardware barrier to training powerful AI models, BAS can help make large-scale AI research more accessible and energy-efficient.