Title: EAGLE: Efficient Analytical Gradient LinearEvaluation for Enhanced Recomputation in Large Language Models
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models; Gradient Checkpointing; Recomputation; Analytical Gradients

Abstract: Training large language models requires substantial memory to store intermediate activations, often exceeding the capacity of modern accelerators. Gradient checkpointing addresses this challenge by trading computation for memory, but introduces significant overhead due to redundant forward passes during recomputation. In this work, we present EAGLE (Efficient Analytical Gradient Linear Evaluation), a novel recomputation strategy that leverages closed-form gradient expressions to eliminate redundant forward computations in linear transformations. Unlike traditional gradient checkpointing that uniformly applies autograd-based recomputation, EAGLE selectively computes gradients analytically for linear layers while maintaining standard recomputation for nonlinear operations. We demonstrate the effectiveness of EAGLE on large-scale models including DeepSeek-236B, DeepSeek-671B, and Llama3-70B, achieving throughput improvements of 1.8\% to 12.5\% compared to full recomputation across models ranging from 70B to 671B parameters. For individual modules, EAGLE delivers substantial recomputation speedups: 6.31× for RMSNorm+Linear, 25\% FLOP reduction for MLP, and up to 2.19× for attention mechanisms, while maintaining identical memory usage and gradient accuracy. The approach offers a triple advantage—memory efficiency matching selective recomputation, computational efficiency approaching no-recomputation baseline, and superior execution time—making it an optimal solution for memory-constrained large-scale model training.

TL;DR: EAGLE is a novel recomputation strategy that analytically computes gradients for linear layers, reducing redundant computation and accelerating large-scale language model training.