Optimal Model Partitioning with Low-Overhead Profiling on the PIM-based Platform for Deep Learning Inference

Seok Young Kim, Jaewook Lee, Yoonah Paik, Chang Hyun Kim, Won Jun Lee, Seon Wook Kim

    Research output: Contribution to journalArticlepeer-review

    Abstract

    Recently Processing-in-Memory (PIM) has become a promising solution to achieve energy-efficient computation in data-intensive applications by placing computation near or inside the memory. In most Deep Learning (DL) frameworks, a user manually partitions a model’s computational graph (CG) onto the computing devices by considering the devices’ capability and the data transfer. The Deep Neural Network (DNN) models become increasingly complex for improving accuracy; thus, it is exceptionally challenging to partition the execution to achieve the best performance, especially on a PIM-based platform requiring frequent offloading of large amounts of data. This article proposes two novel algorithms for DL inference to resolve the challenge: low-overhead profiling and optimal model partitioning. First, we reconstruct CG by considering the devices’ capability to represent all the possible scheduling paths. Second, we develop a profiling algorithm to find the required minimum profiling paths to measure all the node and edge costs of the reconstructed CG. Finally, we devise the model partitioning algorithm to get the optimal minimum execution time using the dynamic programming technique with the profiled data. We evaluated our work by executing the BERT, RoBERTa, and GPT-2 models on the ARM multicores with the PIM-modeled FPGA platform with various sequence lengths. For three computing devices in the platform, i.e., CPU serial/parallel and PIM executions, we could find all the costs only in four profile runs, three for node costs and one for edge costs. Also, our model partitioning algorithm achieved the highest performance in all the experiments over the execution with manually assigned device priority and the state-of-the-art greedy approach.

    Original languageEnglish
    Article number28
    JournalACM Transactions on Design Automation of Electronic Systems
    Volume29
    Issue number2
    DOIs
    Publication statusPublished - 2024 Feb 15

    Bibliographical note

    Publisher Copyright:
    © 2024 Copyright held by the owner/author(s).

    Keywords

    • computational graph
    • optimal scheduling
    • PIM-based execution
    • profiling

    ASJC Scopus subject areas

    • Computer Science Applications
    • Computer Graphics and Computer-Aided Design
    • Electrical and Electronic Engineering

    Fingerprint

    Dive into the research topics of 'Optimal Model Partitioning with Low-Overhead Profiling on the PIM-based Platform for Deep Learning Inference'. Together they form a unique fingerprint.

    Cite this