ABSTRACT As model parameters increase exponentially, distributed training has become essential for advancing modern deep neural networks. Megatron‐LM, an efficient distributed training framework developed by NVIDIA, enables the training of trillion‐parameter models on thousands of GPUs by integrating tensor, pipeline, and data parallelism. Its computational efficiency has established it as a foundational tool for training large‐scale models. Rapid identification of optimal parallel configurations for specific GPU clusters is critical for maximizing computational resource utilization, with training time prediction serving as a key evaluation metric. The high cost and limited availability of high‐performance GPUs, particularly those based on NVIDIA architectures, have made the construction of large‐scale heterogeneous clusters a practical solution to resource and cost constraints. However, existing prediction methods do not reliably or efficiently account for the computational and communication complexities inherent in heterogeneous GPU clusters. To address this gap, HATP (Heterogeneous‐Aware Time Predictor) is introduced as a novel performance prediction method specifically designed for heterogeneous GPU clusters. For any given parallel configuration, HATP rapidly and accurately simulates execution times to inform the optimization of parallel strategies. To address communication differences among heterogeneous GPUs, comprehensive experimental analyses are conducted and analytical expressions are derived to characterize the communication frequency patterns in Megatron‐LM's parallel strategies. This work presents the first systematic quantification of communication operations within Megatron‐LM framework, ensuring that performance predictions remain highly accurate even in complex, heterogeneous environments. Furthermore, to account for computational differences among heterogeneous GPUs, a layer‐level computational performance acquisition scheme is proposed to reduce the impact of fine‐grained operator overlap and additional memory operations. Experimental results demonstrate that HATP achieves an average prediction accuracy of 97.41% in isomorphic environments, surpassing the current state‐of‐the‐art method, ACEso. HATP also attains an average accuracy of 96.04% in heterogeneous data parallel and pipeline parallel configurations, representing the first extension of training time prediction capabilities to heterogeneous environments.
Liu et al. (Thu,) studied this question.