View every non-leaf node as one task. If the current memory usage is \(C\), computing node \(i\) may create a peak of \(C+w_i\). After the task finishes, the children of \(i\) are removed, so the memory usage changes by \(w_i-\sum w_j\), where the sum is over the children of \(i\).

It is convenient to use free memory instead. Suppose the peak memory limit is \(X\), and let \(S\) be the total size of the leaf tensors initially in memory. The initial free memory is \(X-S\). Task \(i\) requires at least \(w_i\) free memory before it starts, and after it finishes the free memory changes by \(\sum w_j-w_i\). Thus each task has a requirement and a gain, and we need an order respecting the tree dependencies that minimizes the necessary initial free memory.

This is exactly the tree scheduling problem solved by Faker’s algorithm (See also: maspy’s Blog). By repeatedly contracting a child block into its parent according to the exchange argument, we can obtain the minimum required initial free memory \(F\). The minimum possible peak memory usage is then \(S+F\). If this value is greater than \(M\), output OOM; otherwise, output the value. The total time complexity is \(O(N\log N)\).