New Tools for Peak Memory Scheduling
CoRR(2023)
摘要
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))$.
更多查看译文
AI 理解论文
溯源树
样例
生成溯源树,研究论文发展脉络
Chat Paper
正在生成论文摘要