New Tools for Peak Memory Scheduling

Abstract

We study scheduling of computation graphs to minimize peak memory consumption, an increasingly critical task due to the surge in popularity of large deep-learning models. This problem corresponds to the weighted version of the classical one-shot black pebbling game. We propose the notion of a dominant schedule to capture the idea of finding the ``best'' schedule for a subgraph and introduce new tools to compute and utilize dominant schedules. Surprisingly, we show that despite the strong requirements, a dominant schedule exists for any computation graph; and, moreover, that it is possible to compute the dominant schedule efficiently whenever we can find optimal schedules efficiently for a particular class of graphs (under mild technical conditions). We apply these new tools to analyze trees and series-parallel graphs. We show that the weighted one-shot black pebbling game is strongly NP-complete even when the graph is an out-tree -- or simpler still, a pumpkin, one of the simplest series-parallel graphs. On the positive side, we design a fixed-parameter tractable algorithm to find a dominant schedule (hence also a peak memory minimizing schedule) for series-parallel graphs when parameterized by the out-degree. This algorithm runs in time $2^{O(d \log d)} \cdot poly(n)$ for series-parallel graphs with $n$ nodes and maximum out-degree $d$; for pumpkins, we can improve the dependence on $d$ to $O(2^d \cdot poly(n))$.