Dynamic Load Balancing for Uncertainty Quantification with Applications in Bayesian Inversion

Abstract

Uncertainty quantification (UQ) workflows present a particular scheduling challenge in high performance computing environments, as they typically generate large numbers of heterogeneous model evaluations with loose but non-trivial dependencies between tasks. A static one-size-fits-all approach in traditional schedulers is inadequate to handle heterogeneous tasks optimally. We introduce an improved load balancer in the UQ and Modelling Bridge (UM-Bridge) framework aimed at mitigating these issues; UM-Bridge is a language-agnostic interface developed to couple UQ software with numerical simulation. This load balancer is first benchmarked against the previous implementation using a synthetic problem. Then, as a realistic example, we test the load balancer with a Bayesian inverse problem solved via multilevel delayed acceptance sampling. The underlying forward problem is a hierarchy of tsunami simulations enabled through ExaHyPE, whose runtimes span several orders of magnitude and loose dependencies between levels make the workload particularly challenging to schedule. Our results indicate the load balancer distributes the sampling requests effectively with an average idle time of close to a millisecond, a thousand times lower than the previous version, without any prior assumptions about the workload or application.