Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Allspark: Workload Orchestration for Visual Transformers on Processing In-Memory Systems

Mengke Ge, Junpeng Wang, Binhan Chen, Yingjian Zhong, Haitao Du, Song Chen, Yi Kang

TL;DR

Allspark is presented, which focuses on workload orchestration for visual Transformers on PIM systems, aiming at minimizing inference latency and speedups on 3D-stacked DRAM-based PIM systems.

Abstract

The advent of Transformers has revolutionized computer vision, offering a powerful alternative to convolutional neural networks (CNNs), especially with the local attention mechanism that excels at capturing local structures within the input and achieve state-of-the-art performance. Processing in-memory (PIM) architecture offers extensive parallelism, low data movement costs, and scalable memory bandwidth, making it a promising solution to accelerate Transformer with memory-intensive operations. However, the crucial issue lies in efficiently deploying an entire model onto resource-limited PIM system while parallelizing each transformer block with potentially many computational branches based on local-attention mechanisms. We present Allspark, which focuses on workload orchestration for visual Transformers on PIM systems, aiming at minimizing inference latency. Firstly, to fully utilize the massive parallelism of PIM, Allspark employs a fine-grained partitioning scheme for computational branches, and formats a systematic layout and interleaved dataflow with maximized data locality and reduced data movement. Secondly, Allspark formulates the scheduling of the complete model on a resource-limited distributed PIM system as an integer linear programming (ILP) problem. Thirdly, as local-global data interactions exhibit complex yet regular dependencies, Allspark provides a two-stage placement method, which simplifies the challenging placement of computational branches on the PIM system into the structured layout and greedy-based binding, to minimize NoC communication costs. Extensive experiments on 3D-stacked DRAM-based PIM systems show that Allspark brings 1.2x-24.0x inference speedup for various visual Transformers over baselines. Compared to Nvidia V100 GPU, Allspark-enriched PIM system yields average speedups of 2.3x and energy savings of 20x-55x.

Allspark: Workload Orchestration for Visual Transformers on Processing In-Memory Systems

TL;DR

Allspark is presented, which focuses on workload orchestration for visual Transformers on PIM systems, aiming at minimizing inference latency and speedups on 3D-stacked DRAM-based PIM systems.

Abstract

The advent of Transformers has revolutionized computer vision, offering a powerful alternative to convolutional neural networks (CNNs), especially with the local attention mechanism that excels at capturing local structures within the input and achieve state-of-the-art performance. Processing in-memory (PIM) architecture offers extensive parallelism, low data movement costs, and scalable memory bandwidth, making it a promising solution to accelerate Transformer with memory-intensive operations. However, the crucial issue lies in efficiently deploying an entire model onto resource-limited PIM system while parallelizing each transformer block with potentially many computational branches based on local-attention mechanisms. We present Allspark, which focuses on workload orchestration for visual Transformers on PIM systems, aiming at minimizing inference latency. Firstly, to fully utilize the massive parallelism of PIM, Allspark employs a fine-grained partitioning scheme for computational branches, and formats a systematic layout and interleaved dataflow with maximized data locality and reduced data movement. Secondly, Allspark formulates the scheduling of the complete model on a resource-limited distributed PIM system as an integer linear programming (ILP) problem. Thirdly, as local-global data interactions exhibit complex yet regular dependencies, Allspark provides a two-stage placement method, which simplifies the challenging placement of computational branches on the PIM system into the structured layout and greedy-based binding, to minimize NoC communication costs. Extensive experiments on 3D-stacked DRAM-based PIM systems show that Allspark brings 1.2x-24.0x inference speedup for various visual Transformers over baselines. Compared to Nvidia V100 GPU, Allspark-enriched PIM system yields average speedups of 2.3x and energy savings of 20x-55x.
Paper Structure (28 sections, 12 equations, 16 figures, 7 tables, 1 algorithm)

This paper contains 28 sections, 12 equations, 16 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Transformer Accelerators
  4. Scheduling and Mapping Space Exploration
  5. Background and Motivation
  6. Visual Transformers
  7. Processing In-Memory Systems
  8. Challenges with End-to-end Inference Deployment
  9. Framework Overview
  10. Partitioning and Dataflow Formation
  11. Key Ideas
  12. NoC-based Dataflow Implementation
  13. Processing Procedures and Weight Burdens per PIM-node
  14. Scheduling for End-to-end Inference
  15. Variables and Constants
  16. ...and 13 more sections

Figures (16)

  • Figure 1: Visual Transformer structure.
  • Figure 2: Shifted window partitioning for local-global interactionSwin.
  • Figure 3: Scalable processing in-memory systems.
  • Figure 4: PIM-node utilization under branch-level parallelism. Swin is deployed to the PIM system with a node array of size 16$\times$16, when the input is $640^2$ and the local region is $7\times 7$. During the execution of the most computationally intensive stage 3 and 4, the utilization is severely below 15%.
  • Figure 5: Allspark overview.
  • ...and 11 more figures