PML-MPI: A Pre-Trained ML Framework for Efficient Collective Algorithm Selection in MPI

The Message Passing Interface is the de facto standard in high-performance computing (HPC) for inter-process communication. MPI libraries employ numerous algorithms for each collective communication pattern whose behavior is largely affected by the underlying hardware, communication pattern, message size, and number of processes involved. Choosing the “best” algorithm for every possible scenario is a non-trivial task. MPI libraries primarily depend on heuristics for algorithm selection on previously unseen clusters, often resulting in evident slowdowns. Although offline micro-benchmarking tools can exhaustively identify optimal algorithms for all configurations, this is an excessively time-consuming approach. Machine Learning (ML) emerged as an alternate approach. However, most ML-based approaches employ online methods that introduce additional runtime overhead, which makes this impractical at scale. To address this challenge, we propose a pre-trained ML framework that eliminates runtime overhead. Our model requires only a quick inference for each new cluster without necessitating model retraining. It incorporates various hardware features to enhance its adaptability across diverse clusters. Our model's training utilizes tuning data from a broad range of architectures, promoting its versatility and our proposed system exhibits up to 6.3% speedup over default heuristics on systems of up to 1024 cores while significantly minimizing model overhead in comparison to existing methodologies.