Robust Scheduling with GFlowNets

Abstract

Finding the best way to schedule operations in a computation graph is a classical NP-hard problem which is central to compiler optimization. However, evaluating the goodness of a schedule on the target hardware can be very time-consuming. Traditional approaches as well as previous machine learning ones typically optimize proxy metrics, which are fast to evaluate but can lead to bad schedules when tested on the target hardware. In this work, we propose a new approach to scheduling by sampling proportionally to the proxy metric using a novel GFlowNet method. We introduce a technique to control the trade-off between diversity and goodness of the proposed schedules at inference time and demonstrate empirically that the pure optimization baselines can lead to subpar performance with respect to our approach when tested on a target model. Furthermore, we show that conditioning the GFlowNet on the computation graph enables generalization to unseen scheduling problems for both synthetic and real-world compiler datasets.

A Theory of Non-acyclic Generative Flow Networks

“…In the midst of the growing literature on GFlowNets (Bengio et al 2021b;Jain et al 2022a;Malkin et al 2022;Madan et al 2022;Zhang et al 2023b;Jain et al 2022b;Pan et al 2023;Deleu and Bengio 2023) three, in particular, are closely related to ours. Bengio et al (2021b) attempt at laying the foundation of the theory of GFlowNets, they discuss many openings for future applications or explorations of the method.…”

“…More sophisticated architectures taking into account the local graph structure and/or learning more global properties of the graph would be desirable. Graph transformer (Dwivedi and Bresson 2020) would be a natural next step, a step already taken on DAG (Zhang et al 2023b) using topoformers (Gagrani et al 2022).…”