Memory and Bandwidth are All You Need for Fully Sharded Data Parallel

TL;DR

This work tackles the challenge of training very large transformers with Fully Sharded Data Parallelism by foregrounding memory and, critically, inter-node bandwidth as bottlenecks. It combines theoretical development, simulation-grid analyses, and extensive experiments across clusters with different network speeds to map how model size, sequence length, and hardware resources interact under FSDP. Key findings show that GPU memory and network bandwidth jointly cap hardware FLOPs utilization (HFU) and model FLOPs utilization (MFU), with throughput (K) bounded by network capabilities, especially for long sequences and very large models. The study provides actionable hardware thresholds and guidelines for configuring FSDP-enabled training to maximize efficiency on bandwidth-constrained clusters, with practical implications for scaling LLMs.

Abstract

Transformer models have revolutionized a wide spectrum of disciplines, especially in language processing. The recent success has proven that model size scalability is crucial for achieving superior performance metrics. However, training large transformer models is challenging even on modern hardware with powerful GPUs and high-speed interconnects. Existing studies primarily focus on optimizing model training distribution strategies to minimize memory footprint and enhance training speed, often overlooking the scalability challenges related to model size and hardware constraints. To address this oversight, we thoroughly investigate computational, memory, and network demands of training large transformers using the Fully Sharded Data Parallel (FSDP) distributed strategy across different hardware clusters. We explore the intricate relationships between model size and hardware setups to identify configurations that ensure maximum model and hardware efficiency, effective sequence length management, and optimal training throughput. A significant finding of our study is the critical interplay of the cluster's connection bandwidth and GPU memory size compared to the computational performance of GPUs. This interplay limits training efficiency, underscoring the role of both hardware characteristics as a possible bottleneck. By integrating theoretical analysis with simulations and empirical tests, we demonstrate how hardware limitations affect training efficacy, identifying key hardware thresholds and the impact of network connectivity. Our findings prompt a reassessment of training strategies guiding users on the way to finding hardware-optimal FSDP configurations, enhancing training efficiency for large-scale transformer models.

Figures (10)

• Figure 1: Representation of theoretical peak MFU and logarithmic throughput (TGS) on 512 GPUs distribution training count across various model sizes. The upper figure presents outcomes from training under Zero stage-3 with activation checkpoints enabled, while the middle represents results from Zero stage-3 without re-computation. The lower panel represent the optimum training strategies derived from exhaustive configuration searches.
• Figure 2: Assessment of MFU and throughput with respect to sequence Length for a 1.3B Model across 4 GPUs. Throughput (TGS) is depicted on a logarithmic scale. The batch size of sequence and context length (ctx in the figure) product, representing batch size in tokens, is utilized as the abscissa.
• Figure 3: Assessment of MFU and throughput with respect to sequence length for a 13B Model across 8 GPUs. Throughput (TGS) is rendered on a logarithmic scale. Notably, for context lengths of 4096 and 8192, tokens per batch are set at 8192, whereas for other configurations, it stands at 10240.
• Figure 4: MFU across different model scales on dual clusters. Models are trained with context lengths optimized for maximum GPU memory usage, employing a batch size of 1. The assessment spans model training utilizing 8 to 512 GPUs. Test outcomes and theoretical maximum MFU predictions vis simulation are presented in each panel.
• Figure 5: Architecture of the transformer block employed in testing.