Distributed Inference with Minimal Off-Chip Traffic for Transformers on Low-Power MCUs

TL;DR

This paper presents a distributed inference methodology for Transformer models on low-power MCUs, significantly reducing off-chip traffic and enabling efficient on-device AI for wearable devices.

Contribution

It introduces a novel distributed system approach that partitions Transformer inference across multiple MCUs, minimizing off-chip traffic and enabling deployment on resource-constrained wearable devices.

Findings

Achieved 26.1x speedup with TinyLlama-42M on 8 MCUs

Reduced energy consumption to 0.64 mJ per inference

Improved Energy Delay Product by 27.2x

Abstract

Contextual Artificial Intelligence (AI) based on emerging Transformer models is predicted to drive the next technology revolution in interactive wearable devices such as new-generation smart glasses. By coupling numerous sensors with small, low-power Micro-Controller Units (MCUs), these devices will enable on-device intelligence and sensor control. A major bottleneck in this class of systems is the small amount of on-chip memory available in the MCUs. In this paper, we propose a methodology to deploy real-world Transformers on low-power wearable devices with minimal off-chip traffic exploiting a distributed system of MCUs, partitioning inference across multiple devices and enabling execution with stationary on-chip weights. We validate the scheme by deploying the TinyLlama-42M decoder-only model on a system of 8 parallel ultra-low-power MCUs. The distributed system achieves an energy consumption of 0.64 mJ, a latency of 0.54 ms per inference, a super-linear speedup of 26.1 x, and an Energy Delay Product (EDP) improvement of 27.2 x, compared to a single-chip system. On MobileBERT, the distributed system's runtime is 38.8 ms, with a super-linear 4.7 x speedup when using 4 MCUs compared to a single-chip system.