A Comparative Analysis of Distributed Training Strategies for GPT-2

TL;DR

The study addresses the high compute and memory costs of training large language models by empirically comparing non-distributed, Fully Sharded Data Parallel (FSDP), and Distributed Data Parallel (DDP) training for GPT-2 on the All the News dataset. It evaluates key metrics such as throughput, memory usage, training time, gradient stability, and loss to understand trade-offs between memory efficiency and speed. Findings show that FSDP excels in memory efficiency, while Single-GPU offers high throughput at the cost of longer training times, and DDP provides a stable middle ground; all methods achieve convergent loss trajectories. These results guide practitioners in selecting appropriate distributed training strategies and motivate future adaptive or hybrid parallelization to scale NLP models further.

Abstract

The rapid advancement in Large Language Models has been met with significant challenges in their training processes, primarily due to their considerable computational and memory demands. This research examines parallelization techniques developed to address these challenges, enabling the efficient and scalable training of Large Language Models. A comprehensive analysis of both data and model parallelism strategies, including Fully Sharded Data Parallelism and Distributed Data-Parallel frameworks, is provided to assess methods that facilitate efficient model training. Furthermore, the architectural complexities and training methodologies of the Generative Pre-Trained Transformer-2 model are explored. The application of these strategies is further investigated, which is crucial in managing the substantial computational and memory demands of training sophisticated models. This analysis not only highlights the effectiveness of these parallel training strategies in enhancing training efficiency but also their role in enabling the scalable training of large language models. Drawing on recent research findings, through a comprehensive literature review, this research underscores the critical role of parallelization techniques in addressing the computational challenges of training state-of-the-art Large Language Models, thereby contributing to the advancement of training more sophisticated and capable artificial intelligence systems.

Figures (6)

• Figure 1: GPT-2 Architecture
• Figure 2: Throughput per Epoch
• Figure 3: Grad Norm per Epoch
• Figure 4: Training Time per Epoch
• Figure 5: Memory usage per epoch