A structure-aware framework for learning device placements on computation graphs

Overview

• Existing approaches for device placement ignore the topological features of computation graphs and rely mostly on heuristic methods for graph partitioning.
• They either follow a grouper-placer or an encoder-placer architecture, which requires understanding the interaction structure between code operations.
• This paper proposes a novel framework for device placement, using smaller computation graphs extracted from the OpenVINO toolkit and reinforcement learning.
• The framework consists of five steps, including graph coarsening, node representation learning, and policy optimization.
• It facilitates end-to-end training and considers the directed and acyclic nature of the computation graphs.
• The paper also proposes a model variant inspired by graph parsing networks and complex network analysis, enabling joint graph representation learning and personalized graph partitioning.

Plain English Explanation

When it comes to device placement, existing approaches often ignore the inherent structure and connections within the computational graphs, relying instead on basic heuristic methods for dividing the graph. These techniques either follow a "grouper-placer" or an "encoder-placer" architecture, which requires understanding the intricate relationships between different operations in the code.

To address this, the researchers in this paper propose a novel framework that uses smaller computational graphs extracted from the OpenVINO toolkit and reinforcement learning to determine the optimal device placement. The framework goes through five key steps, including simplifying the graph (graph coarsening), learning how to represent the nodes (node representation learning), and optimizing the placement policy. This end-to-end approach takes into account the directed and acyclic nature of the computational graphs, which is an important consideration.

The paper also introduces a model variant inspired by graph parsing networks and complex network analysis. This variant can simultaneously learn how to represent the graph and partition it into an unspecified number of groups, tailoring the placement to the specific computational graph.

Technical Explanation

The key aspects of this paper's technical approach are:

1. Graph Coarsening: The researchers start by extracting smaller computation graphs from the larger OpenVINO toolkit using a graph coarsening technique. This simplifies the problem and makes it more manageable for the reinforcement learning-based optimization.

2. Node Representation Learning: Next, they learn numerical representations for each node in the simplified computation graphs. This allows the model to capture the important features and relationships within the graph structure.

3. Policy Optimization: The core of the framework is a reinforcement learning-based optimization process that learns the optimal device placement policy. The model is trained to maximize the execution speed of the suggested placements.

4. Graph Parsing Variant: The paper also introduces a model variant that combines graph representation learning and personalized graph partitioning. This enables the framework to adaptively partition the computation graph into an unspecified number of groups, tailoring the placement to the specific graph structure.

The researchers evaluate their approach on three benchmark models: Inception-V3, ResNet, and BERT. They demonstrate significant improvements in inference speed, with up to 58.2% faster execution compared to CPU-only and up to 60.24% faster than other baselines.

Critical Analysis

The researchers have proposed an innovative approach that addresses the limitations of existing device placement techniques. By leveraging reinforcement learning and considering the topological structure of computation graphs, the framework can generate more efficient placements.

However, the paper does not provide a detailed discussion of the limitations or potential issues with the proposed methods. For example, the computational complexity of the graph coarsening and node representation learning steps could be a concern, especially for larger and more complex graphs. Additionally, the reliance on the OpenVINO toolkit may limit the generalizability of the approach to other frameworks or custom computational graphs.

Further research could explore the scalability of the framework, its performance on a wider range of models and hardware configurations, and the potential trade-offs between placement quality and computational efficiency. Incorporating additional constraints, such as power consumption or memory usage, could also enhance the practical applicability of the framework.

Conclusion

This paper presents a novel device placement framework that addresses the shortcomings of existing approaches by leveraging reinforcement learning and incorporating the topological features of computation graphs. The proposed solution demonstrates significant improvements in inference speed across various benchmark models, showcasing the potential of this data-driven approach to optimizing hardware utilization.

The paper's innovative techniques, such as the graph coarsening, node representation learning, and the graph parsing variant, provide a solid foundation for further advancements in the field of automated hardware-software co-design and constrained object placement. As the complexity of computational workloads continues to grow, this research highlights the importance of considering the structural properties of computation graphs in developing efficient device placement strategies.