Hyperparameter optimization of data-driven AI models on HPC systems

TL;DR

This paper demonstrates large-scale hyperparameter optimization of AI models on HPC systems, showing significant performance improvements using advanced algorithms like ASHA combined with Bayesian optimization.

Contribution

It benchmarks and compares hyperparameter search algorithms on HPC resources, highlighting the effectiveness of ASHA with Bayesian optimization for large-scale AI model tuning.

Findings

Hyperparameter optimization improved model performance significantly.

ASHA with Bayesian optimization was most resource-efficient.

Large-scale HPC resources enabled optimization impossible otherwise.

Abstract

In the European Center of Excellence in Exascale computing "Research on AI- and Simulation-Based Engineering at Exascale" (CoE RAISE), researchers develop novel, scalable AI technologies towards Exascale. This work exercises High Performance Computing resources to perform large-scale hyperparameter optimization using distributed training on multiple compute nodes. This is part of RAISE's work on data-driven use cases which leverages AI- and HPC cross-methods developed within the project. In response to the demand for parallelizable and resource efficient hyperparameter optimization methods, advanced hyperparameter search algorithms are benchmarked and compared. The evaluated algorithms, including Random Search, Hyperband and ASHA, are tested and compared in terms of both accuracy and accuracy per compute resources spent. As an example use case, a graph neural network model known as MLPF, developed for the task of Machine-Learned Particle-Flow reconstruction in High Energy Physics, acts as the base model for optimization. Results show that hyperparameter optimization significantly increased the performance of MLPF and that this would not have been possible without access to large-scale High Performance Computing resources. It is also shown that, in the case of MLPF, the ASHA algorithm in combination with Bayesian optimization gives the largest performance increase per compute resources spent out of the investigated algorithms.