Stochastic Rounding, Not Momentum: Full Fine-Tuning at 2.25 Bytes per Parameter
Authors/Creators
Description
Full fine-tuning with AdamW stores about 16 bytes per parameter before activations, so a 7B model does not fit a 24 GB GPU. We show that a much simpler update suffices: each tensor is updated as soon as its gradient is complete, by the gradient divided by a decaying maximum of its magnitude over each block of 16 weights, with no momentum, and the result is written to bf16 weights with stochastic rounding, so training holds only the weights plus one fp32 scale per 16 weights (2.25 bytes per parameter). On code fine-tuning of Qwen2.5-Coder-1.5B and 3B and SmolLM2-1.7B, with every method at its own tuned learning rate and three seeds, it matches or beats FlashAdamW and AdaLomo, lowering held-out loss by 0.0036 relative to FlashAdamW on SmolLM2 with 44% less peak memory, and it fully fine-tunes Qwen2.5-Coder-7B on one 24 GB GPU. Two findings explain why it works: stochastic rounding is necessary, since every variant that rounds low-precision state to nearest loses part of the small streamed steps (and adding it to AdaLomo closes most of AdaLomo's gap to us), and momentum is unnecessary once updates are streamed and the learning rate decays to zero, while the decaying block maximum is not. Code: https://github.com/marchfederico/blockmax-finetune
Files
arxiv_preview.pdf
Additional details
Related works
- Is supplemented by
- Software: https://github.com/marchfederico/blockmax-finetune (URL)
Software
- Repository URL
- https://github.com/marchfederico/blockmax-finetune