vix.ing · top · new · best · stats · spec

New Tools for Peak Memory Scheduling

2023/12/21 by Ce Jin, Manish Purohit, Jin, Ce +7
Computer Science · #Complexity and Algorithms in Graphs #Data Structures and Algorithms (cs.DS) #FOS: Computer and information sciences #Graph Theory and Algorithms #Parallel Computing and Optimization Techniques

paper · pdf · doi:10.48550/arxiv.2312.13526

openalex publication_date 2023/12/21 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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 2O(d log d) ⋅ 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(2d ⋅ poly(n)).

Related