Efficient Algorithms for Device Placement of DNN Graph Operators

Abstract

Modern machine learning workloads use large models, with complex structures, that are very expensive to execute. The devices that execute complex models are becoming increasingly heterogeneous as we see a flourishing of domain-specific accelerators being offered as hardware accelerators in addition to CPUs. These trends necessitate distributing the workload across multiple devices. Recent work has shown that significant gains can be obtained with model parallelism, i.e, partitioning a neural network's computational graph onto multiple devices. In particular, this form of parallelism assumes a pipeline of devices, which is fed a stream of samples and yields high throughput for training and inference of DNNs. However, for such settings (large models and multiple heterogeneous devices), we require automated algorithms and toolchains that can partition the ML workload across devices. In this paper, we identify and isolate the structured optimization problem at the core of device placement of DNN operators, for both inference and training, especially in modern pipelined settings. We then provide algorithms that solve this problem to optimality. We demonstrate the applicability and efficiency of our approaches using several contemporary DNN computation graphs.

Software-hardware co-design for fast and scalable training of deep learning recommendation models

“…FlexFlow [22] uses automatic search to discover the best operator parallelization strategy in the graph. Building on this direction of auto-parallelization, these recent papers [39,60] use optimal synthesis and reinforcement learning to find optimized device placement to further improve parallelism without the need for manual intervention. However, these general systems are not specifically designed for highly sparse recommendation models.…”