Pipette: Automatic Fine-grained Large Language Model Training Configurator for Real-World Clusters

TL;DR

Pipette addresses the challenge of automatically configuring fine-grained 3D parallelism for real-world LLM training by profiling heterogeneous interconnects, incorporating a memory-aware estimator, and mapping GPUs with simulated-annealing-based fine-grained dedication. It introduces a refined latency model that captures hidden paths in memory-efficient 1F1B schedules and a neural memory estimator to prevent out-of-memory configurations, enabling practical auto-configuration. The approach yields up to ~1.46× speedup over AMP and ~1.56× over Varuna on large clusters, while reducing memory errors and keeping configuration overhead negligible relative to training time. The work provides a practical, scalable, and open-source solution for deploying efficient LLM training on diverse real-world clusters.

Abstract

Training large language models (LLMs) is known to be challenging because of the huge computational and memory capacity requirements. To address these issues, it is common to use a cluster of GPUs with 3D parallelism, which splits a model along the data batch, pipeline stage, and intra-layer tensor dimensions. However, the use of 3D parallelism produces the additional challenge of finding the optimal number of ways on each dimension and mapping the split models onto the GPUs. Several previous studies have attempted to automatically find the optimal configuration, but many of these lacked several important aspects. For instance, the heterogeneous nature of the interconnect speeds is often ignored. While the peak bandwidths for the interconnects are usually made equal, the actual attained bandwidth varies per link in real-world clusters. Combined with the critical path modeling that does not properly consider the communication, they easily fall into sub-optimal configurations. In addition, they often fail to consider the memory requirement per GPU, often recommending solutions that could not be executed. To address these challenges, we propose Pipette, which is an automatic fine-grained LLM training configurator for real-world clusters. By devising better performance models along with the memory estimator and fine-grained individual GPU assignment, Pipette achieves faster configurations that satisfy the memory constraints. We evaluated Pipette on large clusters to show that it provides a significant speedup over the prior art. The implementation of Pipette is available at https://github.com/yimjinkyu1/date2024_pipette.

Figures (9)

• Figure 1: Example 3D parallelism configuration with 12 GPUs.
• Figure 2: Pipeline scheduling.
• Figure 3: Inter-stage communication latency in a real-world cluster for 40 days.
• Figure 4: Baseline and fine-grained worker dedication schedule.
• Figure 5: Latency estimation and top 10 recommendation of baselines and Pipette.