Taming LLMs by Scaling Learning Rates with Gradient Grouping

TL;DR

This paper introduces Scaling with Gradient Grouping (SGG), a novel optimizer wrapper that enhances adaptive learning rate estimation for large language models by dynamically grouping gradients, leading to more stable, faster, and more compatible training.

Contribution

The paper proposes SGG, a new method that improves learning rate estimation through gradient grouping and scaling, enhancing LLM training stability and efficiency.

Findings

SGG improves training stability across different batch sizes.

SGG accelerates convergence compared to baseline optimizers.

SGG demonstrates consistent performance gains on diverse LLM benchmarks.

Abstract

Training large language models (LLMs) poses challenges due to their massive scale and heterogeneous architectures. While adaptive optimizers like AdamW help address gradient variations, they still struggle with efficient and effective parameter-wise learning rate estimation, resulting in training instability, slow convergence, and poor compatibility with parameter-efficient fine-tuning (PEFT) techniques. This work introduces Scaling with Gradient Grouping (SGG), an optimizer wrapper that improves adaptive learning rate estimation by dynamic grouping and group-specific scaling. SGG first groups gradient statistics in each layer into clusters and then applies cluster-specific scaling to calibrate learning rates for each parameter, thus imposing collective group-wise constraints while maintaining precise per-parameter adaptation. Experiments on diverse (M)LLM benchmarks show that SGG integrates seamlessly with existing optimizers, and offers consistent gains and faster convergence over baselines, with various model sizes. Its stability across varying batch sizes and learning rates establishes SGG as a robust choice for LLM optimization.