Graph rewriting for large model support using categorized topological sort

A computer program product for managing GPU memory consumption by computational graph rewriting, the computer program product comprising a non-transitory computer readable storage medium having program instructions embodied therewith, the program instructions executable by a computer to cause the computer to perform a method comprising:

• constructing, by a hardware processor, a categorized topological ordering of a computational graph, the categorized topological ordering including multiple computational nodes arranged in multiple levels;
• estimating, by the hardware processor, the GPU memory consumption responsive to a level including two or more computational nodes from among the multiple computational nodes;
• rewriting, by the hardware processor, the computational graph by linearizing the two or more computational nodes in the level to avoid overlapping of the GPU memory consumption by the two or more computational nodes responsive to the GPU memory consumption exceeding a threshold;
• and
• managing the GPU memory consumption in accordance with the rewritten computational graph.

A computer processing system for managing GPU memory consumption by computational graph rewriting, comprising:

• a memory for storing program code;
• and
• a hardware processor, operatively coupled to the memory, for running the program code to
• construct a categorized topological ordering of a computational graph, the categorized topological ordering including multiple computational nodes arranged in multiple levels;
• estimate the GPU memory consumption responsive to a level including two or more computational nodes from among the multiple computational nodes;
• rewrite the computational graph by linearizing the two or more computational nodes in the level to avoid overlapping of the GPU memory consumption by the two or more computational nodes responsive to the GPU memory consumption exceeding a threshold;
• and
• manage the GPU memory consumption in accordance with the rewritten computational graph.