GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling
Tianhao Chen, Xin Xu, Zijing Liu, Pengxiang Li, Xinyuan Song, AJAY JAISWAL, Fan Zhang, Jishan Hu, Yang Wang, Hao CHEN, Shizhe Diao, Shiwei Liu, Yu Li, Lu Yin, Can Yang
Modern Large Language Models, such as the LLaMA, Qwen and DeepSeek series, predominantly adopt the Pre-LayerNorm (Pre-LN) transformer architecture. While Pre-LN is stable during pretraining and scales well to large model sizes, it can suffer from an exponential growth in activation variance across layers. A straightforward analysis of the layer-wise Jacobian, supported by empirical evidence, shows that this phenomenon causes deeper layers to contribute less to learning, thereby slowing the convergence speed and limiting model performance. To mitigate this issue, we propose Gradient-Preserving Activation Scaling (GPAS), a simple technique that can be used in combination with existing approaches. GPAS works by scaling down the intermediate activations while keeping their gradients unchanged. This leaves information in the activations intact, and avoids the gradient vanishing problem associated with gradient downscaling. Extensive experiments across various model sizes from 71M to 1B show that GPAS achieves consistent performance gain. Beyond enhancing Pre-LN transformers, GPAS also shows promise in improving alternative architectures such as Sandwich-LN and Post-LN, demonstrating its versatility and potential for improving training dynamics in a wide range of settings.