Differentiable Combinatorial Scheduling at Scale

TL;DR

This work addresses resource-constrained scheduling, an NP-hard problem, by introducing a differentiable combinatorial scheduling framework built on System of Difference Constraints (SDC) and a constrained Gumbel Trick. The method vectorizes the search space, encodes inequality constraints via a differentiable transformation, and optimizes a dual-objective loss that minimizes memory and inter-stage communication using gradient descent on GPUs. Key contributions include a fully differentiable LP-based scheduling formulation without training data and demonstrated performance gains over CPLEX, Gurobi, and CP-SAT on synthetic and EPFL benchmarks. The approach broadens the applicability of differentiable methods to large-scale scheduling problems and enables fast, parallelizable optimization for DAG-based workloads, with open-source code available at the project repository.

Abstract

This paper addresses the complex issue of resource-constrained scheduling, an NP-hard problem that spans critical areas including chip design and high-performance computing. Traditional scheduling methods often stumble over scalability and applicability challenges. We propose a novel approach using a differentiable combinatorial scheduling framework, utilizing Gumbel-Softmax differentiable sampling technique. This new technical allows for a fully differentiable formulation of linear programming (LP) based scheduling, extending its application to a broader range of LP formulations. To encode inequality constraints for scheduling tasks, we introduce \textit{constrained Gumbel Trick}, which adeptly encodes arbitrary inequality constraints. Consequently, our method facilitates an efficient and scalable scheduling via gradient descent without the need for training data. Comparative evaluations on both synthetic and real-world benchmarks highlight our capability to significantly improve the optimization efficiency of scheduling, surpassing state-of-the-art solutions offered by commercial and open-source solvers such as CPLEX, Gurobi, and CP-SAT in the majority of the designs.

Figures (5)

• Figure 1: Example of SDC-based scheduling --- (left) A DFG with two schedule stages $l_{0}$ and $l_{1}$ with latency $L=3$; (right) Dependence constraints and objective functions including peak memory minimization and inter-stage communication minimization (the blue crosses)
• Figure 2: The implementation overview of our differentiable scheduling. $p$ indicates the probability vector for each node during GS-enabled differentiable optimization; $T$ indicates the dependency constraints from the predecessors $S$ indicates the scheduled stage for the node. (a) The search space vectorization and GS initialization. (b) The legal scheduling after applying the dependency constraints. (c) The scheduling is optimized with GS-enabled differentiable optimization.
• Figure 3: Performance comparisons with random workloads. Baseline results are SDC scheduling solved by commercial SOTA CPLEX, Gurobi, and CP-SAT solvers.
• Figure 4: Performance comparisons with workloads built for EPFL benchmarks. Baseline results are SDC scheduling solved by commercial SOTA CPLEX, Gurobi, and CP-SAT solvers.
• Figure 5: Loss Function Evaluation with selected Random workloads and EPFL designs.

Theorems & Definitions (2)

• Lemma 3.1
• proof