Optimizing Large Model Training through Overlapped Activation Recomputation
This work addresses efficiency bottlenecks in training large neural networks, offering incremental improvements for researchers and practitioners in deep learning.
The paper tackles the high overhead of activation recomputation in large model training by introducing Lynx, a framework that overlaps recomputation with communication and uses heuristic scheduling and model partitioning, achieving up to 1.37x speedup over existing methods in GPT models with 1.3B-23B parameters.
Large model training often uses recomputation to alleviate memory pressure and pipelines to exploit the parallelism of data, tensors, and devices. However, existing recomputation approaches may incur high overhead when training real-world models, as they are executed on demand in the critical training path. In this paper, we present Lynx, a new recomputation framework to reduce overhead by overlapping recomputation with communication in training pipelines. To reduce the large search space for recomputation strategies, we propose a heuristic-based recomputation scheduling algorithm, which is based on the observation that there are identical structures in large DNN models so that we can apply the same scheduling policy to all such structures. Additionally, we propose a recomputation-aware model partitioning method to balance each stage's execution time for improved training throughput. Our comprehensive evaluation using GPT models with 1.3B-23B parameters shows that Lynx outperforms existing recomputation approaches by up to 1.37x.