小批量训练可稳定高效,无需梯度累积
Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful
- 按token数固定二阶矩衰减半衰期,适配小批量
- 批量为1时仍稳定,且每计算量性能更优
- 适合资源有限场景,可替代LoRA节省显存
传统观点认为小批量训练语言模型不稳定,因此常采用梯度累积以提升有效批次大小。本文将批量缩小至1,并提出调整Adam超参数的新规则:不固定二阶矩衰减率,而是固定其在训练令牌数上的半衰期。实验发现,小批量(包括批量为1)训练稳定,对超参数更鲁棒,且单位计算量性能优于大批次。更重要的是,即使使用无动量的普通SGD,无需存储优化器状态,也能实现稳定训练。基于此,我们建议除非在多设备多副本训练时,否则应避免梯度累积。此外,小批量配合小状态优化器可达到全微调性能,同时内存占用与LoRA相当。
原文摘要 · Abstract (English)
Conventional wisdom dictates that small batch sizes make language model pretraining and fine-tuning unstable, motivating gradient accumulation, which trades off the number of optimizer steps for a proportional increase in batch size. While it is common to decrease the learning rate for smaller batch sizes, other hyperparameters are often held fixed. In this work, we revisit small batch sizes all the way down to batch size one, and we propose a rule for scaling Adam hyperparameters to small batch sizes. In particular, rather than holding the decay rate of the second moment fixed across batch sizes, we propose to hold its half-life fixed in terms of tokens. We find that small batch sizes (1) train stably, (2) are consistently more robust to hyperparameter choices, (3) achieve equal or better per-FLOP performance than larger batch sizes, and (4) notably enable stable language model training with vanilla SGD, even without momentum, despite storing no optimizer state. Building on these results, we provide practical recommendations for selecting a batch size and setting optimizer hyperparameters. We further recommend against gradient accumulation unless training on multiple devices with multiple model replicas. Finally, we show that a small batch size combined with an optimizer with a small state size can provide the performance benefits of full fine-tuning while maintaining a similar memory footprint to LoRA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。