Practical Machine Learning Autotuning for Large-Scale Collective Communication

Abstract

Collective communication is a fundamental communication model for parallel computing on distributed memory systems. The performance of a collective operation depends on the underlying algorithm. Machine learning (ML)–based autotuners can optimize algorithm selection to enhance collective performance. Previous approaches are impractical for production use at large scale, however, due to prohibitively long training times that exceed typical job durations. This paper introduces ACCLAiM (Advancing Collective Communication Autotuning using Machine Learning), the first ML-based collective algorithm selection autotuner capable of accelerating production applications on large-scale supercomputers. ACCLAiM incorporates several improvements over prior designs in training point selection, handling of non-power-of-two feature values, model validation, and data collection. Our approach leverages variance-based active learning alongside topology-aware benchmark parallelization to eliminate unnecessary training points and maximize machine utilization, thereby significantly reducing training time. We present ACCLAiM as an open-source prototype and provide a comprehensive experimental evaluation. We demonstrate how each of ACCLAiM’s enhancements contributes to a substantial reduction in training time compared with the previous state-of-the-art approach, cumulatively reducing the training time to 5–10 minutes even at large scale. We demonstrate ACCLAiM’s usefulness on two leadership-class supercomputers and showcase its practical benefits for applications, achieving speedups up to 4.1x.