The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training

TL;DR

This paper identifies a sharpness disparity among transformer blocks during training and introduces a blockwise learning rate strategy that accelerates language model pre-training by nearly two times while reducing memory usage.

Contribution

It uncovers the sharpness disparity in transformer blocks and proposes a blockwise learning rate method that improves training speed and efficiency for large language models.

Findings

Achieves nearly 2x speedup in LLM pre-training.

Reduces memory usage by 2x with the new method.

Demonstrates effectiveness across multiple models and datasets.

Abstract

Transformers consist of diverse building blocks, such as embedding layers, normalization layers, self-attention mechanisms, and point-wise feedforward networks. Thus, understanding the differences and interactions among these blocks is important. In this paper, we uncover a clear Sharpness Disparity across these blocks, which emerges early in training and intriguingly persists throughout the training process. Motivated by this finding, we propose Blockwise Learning Rate (LR), a strategy that tailors the LR to each block's sharpness, accelerating large language model (LLM) pre-training. By integrating Blockwise LR into AdamW, we consistently achieve lower terminal loss and nearly $2\times$ speedup compared to vanilla AdamW. We demonstrate this acceleration across GPT-2 and LLaMA, with model sizes ranging from 0.12B to 2B and datasets of OpenWebText, MiniPile, and C4. Finally, we incorporate Blockwise LR into Adam-mini (Zhang et al., 2024), a recently proposed memory-efficient variant of Adam, achieving a combined $2\times$ speedup and $2\times$ memory saving. These results underscore the potential of exploiting the sharpness disparity to improve LLM training.