Robust Scheduling with GFlowNets

要約

計算グラフで操作をスケジュールする最適な方法を見つけることは、コンパイラの最適化の中心となる古典的な NP 困難な問題です。
ただし、ターゲット ハードウェアのスケジュールの良さを評価するには、非常に時間がかかる場合があります。
従来のアプローチと以前の機械学習のアプローチは、通常、プロキシ メトリックを最適化します。これは、評価が高速ですが、ターゲット ハードウェアでテストしたときに不適切なスケジュールにつながる可能性があります。
この作業では、新しい GFlowNet メソッドを使用してプロキシ メトリックに比例してサンプリングすることにより、スケジューリングへの新しいアプローチを提案します。
推論時に提案されたスケジュールの多様性と良さの間のトレードオフを制御する手法を導入し、ターゲット モデルでテストした場合、純粋な最適化ベースラインがアプローチに関して標準以下のパフォーマンスにつながる可能性があることを経験的に示します。
さらに、計算グラフ上で GFlowNet を条件付けすることで、合成および実世界のコンパイラ データセットの両方について、目に見えないスケジューリングの問題を一般化できることを示します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google