Wainscot: Tailoring Model Parallelism to Fit Device Memory Limits

With increasing sizes of DNN (Deep Neural Network) models making them exceed the memory of a single device (GPU), model parallelism-based training has become paramount, splitting a model across multiple devices. Unfortunately, today’s model parallelism approaches often result in memory-unbalanced allocations of the model across the multiple devices, with some devices’ memory heavily utilized while others remain underused. This imbalance limits deployments from reaching high batch sizes, triggers Out of Memory (OOM) errors earlier, and underutilizes resources. We present Wainscot, a model parallelism solution that produces memory-balanced placements of a DNN model across multiple devices, without noticeable increases in step time. We explore and empirically compare different granularities of rebalancing: operators, operator groups, and subgraphs. Experiments with diverse DNNs across a wide range of batch sizes demonstrate that compared to state-of-the-art model parallelism systems, Wainscot reduces maximal peak memory (across devices) by 47.94%, with a modest increase of 1.62% in step time.