Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving inference time in multi-TPU systems with profiled model segmentation

Jorge Villarrubia, Luis Costero, Francisco D. Igual, Katzalin Olcoz

TL;DR

This work address the memory-bound bottleneck of Edge TPU inference caused by limited on-chip memory and host memory transfers. It combines profiling-guided model segmentation with pipelined multi-TPU execution to distribute weights and computation across up to four devices. The study shows substantial gains, including up to 46x speedups for fully connected and up to 6x for convolutional models over a single TPU, demonstrating the practical impact of memory-aware partitioning. These results guide deployment of edge accelerators by balancing device memory usage and pipeline latency, with implications for energy efficiency and hybrid CPU-TPU workflows in future work.

Abstract

In this paper, we systematically evaluate the inference performance of the Edge TPU by Google for neural networks with different characteristics. Specifically, we determine that, given the limited amount of on-chip memory on the Edge TPU, accesses to external (host) memory rapidly become an important performance bottleneck. We demonstrate how multiple devices can be jointly used to alleviate the bottleneck introduced by accessing the host memory. We propose a solution combining model segmentation and pipelining on up to four TPUs, with remarkable performance improvements that range from $6\times$ for neural networks with convolutional layers to $46\times$ for fully connected layers, compared with single-TPU setups.

Improving inference time in multi-TPU systems with profiled model segmentation

TL;DR

This work address the memory-bound bottleneck of Edge TPU inference caused by limited on-chip memory and host memory transfers. It combines profiling-guided model segmentation with pipelined multi-TPU execution to distribute weights and computation across up to four devices. The study shows substantial gains, including up to 46x speedups for fully connected and up to 6x for convolutional models over a single TPU, demonstrating the practical impact of memory-aware partitioning. These results guide deployment of edge accelerators by balancing device memory usage and pipeline latency, with implications for energy efficiency and hybrid CPU-TPU workflows in future work.

Abstract

In this paper, we systematically evaluate the inference performance of the Edge TPU by Google for neural networks with different characteristics. Specifically, we determine that, given the limited amount of on-chip memory on the Edge TPU, accesses to external (host) memory rapidly become an important performance bottleneck. We demonstrate how multiple devices can be jointly used to alleviate the bottleneck introduced by accessing the host memory. We propose a solution combining model segmentation and pipelining on up to four TPUs, with remarkable performance improvements that range from for neural networks with convolutional layers to for fully connected layers, compared with single-TPU setups.

Paper Structure

This paper contains 13 sections, 6 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Background
  3. The Edge TPU architecture
  4. Related work
  5. Inference performance on a single TPU
  6. Synthetic model generation
  7. Performance evaluation
  8. Analysis of host and device memory usage
  9. Model segmentation on multiple TPUs
  10. Performance on single inputs with segmented models
  11. Performance on batched inputs with segmented models
  12. Optimizing model segmentation with profiling
  13. Conclusions

Figures (6)

  • Figure 1: Example of a 3x3 systolic matrix and the cycle-by-cycle data flow through the chains.
  • Figure 2: Inference time analysis for a single Edge TPU.
  • Figure 3: Top: Single-TPU execution for a segmented model with layers stored in the host and device memory. Bottom: Implemented pipelined execution scheme.
  • Figure 4: Inference time for fc models (top) and conv models (bottom) on multiple TPUs using a pipelined implementation.
  • Figure 5: Inference time for segmented models with profiling for fc and conv layers using multiple TPUs for a 50-input batch.
  • ...and 1 more figures