Hermes: Memory-Efficient Pipeline Inference for Large Models on Edge Devices
^††thanks: This work was supported by Shanghai Key Laboratory of Scalable Computing and Systems, National Key Laboratory of Ship Structural Safety, and the Eighth Research Institute of China Aerospace Science and Technology Group Company, Ltd., under Grant USCAST2023-17 and Grant USCAST2023-21. (Corresponding author: Ruhui Ma.)

The architectural makeup of transformer models is the basis for developing pipeline inference strategies on edge devices. A typical transformer model comprises embedding, encoder, decoder, pooling, and additional specialized layers. Based on their architectural configurations, they can be classified into three primary categories: encoder-decoder, encoder-only and decoder-only models. Encoder-decoder models, such as BART [32] and T5 [33] integrate both encoder and decoder layers. BART is architected for complex sequence-to-sequence tasks, and T5 generalizes this capability with a comprehensive text-to-text methodology applicable to a wide array of NLP challenges. Encoder-only models like BERT [34] employ a series of encoder layers to interpret input data, ideally suited for tasks that do not involve generating sequences. Meanwhile, Vision Transformer (ViT) presents an innovative adaptation of encoder-only architecture to analyze sequences of image patches, eschewing the conventional decoder layout found in text-centric models. Conversely, decoder-only models such as GPT rely exclusively on decoder layers, focusing on the generation of new content by drawing on recognizable patterns.

In order to design an efficient pipeline scheme, we conduct experiments to characterize two key aspects when the model performs forward computation: the allocation of memory among transformer model layers and the latency of model loading and model inference. We first evaluate the memory allocation in five kinds of transformer models, including ViT-Large, BERT-Large, GPT-2, GPT-J and BART (BART-Base and BART-Large), which cover all three categories of Transformer models. Additionally, we evaluate the time requirements of loading and inference for various transformer models, including BERT-Large, GPT-2, ViT-Large and GPT-J, by performing standard model inference. All the experiments are conducted on Intel(R) 193 Xeon(R) Gold 6248R CPU.

[M]Observation I For general transformer-based models, the encoder or decoder layers occupy the largest memory footprint. TestCase2: Latency evaluation. To evaluate the latency of model loading and inference, we run standard model inference processes for four transformer models on CPU. Generally, transformer models exhibit considerably higher latency during layer loading compared to layer inference. Through our experiments, as depicted in Fig. 3, we observe that, for the first three smaller models (each with a memory footprint around 1 GB), the layer loading period substantially exceeds the inference time, by roughly an order of magnitude. Conversely, for the larger GPT-J model (12 GB), the layer loading duration is approximately twice that of the inference time. Subsequently, such disparities contribute to a significant portion of the computational process, between 60% to 80%, being spent idle during typical pipeline execution, underlining a serious pipeline stall issue, as shown in Fig. 1(b).

[M]Observation II For general transformer-based models, loading latency is much larger than the inference latency, resulting in the execution process being stalled during most of inference time.

Our experiments in §II-B analyze the time distribution and memory usage of transformer-based large models during model inference. Observation I suggests that a targeted focus on either the encoder or decoder layers is pivotal for optimizing memory management in our pipeline infrastructure. Observation II underscores the necessity of adopting a parallel loading strategy by overlapping multiple inference times with a single loading time within our pipeline scheme, to efficiently mitigate pipeline stalls. In summary, we progress our design by addressing the following challenges: (1) memory challenge on edge devices; (2) pipeline stall problem caused by the huge gap between loading and inference latency.

We present PipeLoad, a memory-efficient pipeline execution mechanism to reduce memory footprint and latency during model inference on edge devices. There are three core workers in PipeLoad mechanism: multiple Loading Agents, one Inference Agent and one Daemon Agent. Loading Agents work in parallel to load model layers from disk to memory, reducing inference latency. The Inference Agent simultaneously executes computations on these loaded layers sequentially in CPU, guaranteeing the model’s predictive accuracy and minimizing pipeline stalls. The Daemon Agent maintains a queue of loaded layers in memory, detects memory usage and destroys memory space for specific layers at a specific point to reduce memory overhead. Three workers communicate with each other through a signalling mechanism that facilitates the realization of whole memory-efficient pipeline.

Fig. 4 illustrates the overall workflow of PipeLoad, including the loading and inference process of model layers as well as the signaling mechanism. Before performing pipeline inference, we adopt a layer-based model partitioning scheme to pre-process the model weights. Multiple Loading Agents constantly load specific layers from disk into memory in parallel. Once a model layer is successfully loaded, the corresponding Loading Agent transmits a computation ready signal that corresponds to this layer to Inference Agent, indicating that this layer is ready for computation. Inference Agent maintains an inference queue in CPU that decides which layer will be processed next, ensuring that model inference respects the original sequence of layers. Upon receiving the computation ready signal, Inference Agent performs forward computation only if all preceding layers have been computed. Following computation of the layer, Inference Agent issues a memory destruction signal to Daemon Agent, notifying it to destroy the memory space of the layer. Daemon Agent then destroys the memory space occupied by the layer to reserve enough space for other layers. When memory usage is about to exceed or has exceeded the memory constraints of the edge device, Daemon Agent sends a stop signal to all Loading Agents, pausing their loading operations until sufficient memory space is available.

Fig. 5 presents a simple case of PipeLoad with three Loading Agents. Based on the characteristics of transformer model layers, we adopt a layer-based model partitioning scheme. In our scheme, we methodically segment the general transformer model architecture into its constituent layers: embedding layers, encoder layers, decoder layers and other layers. Among these layers, we focus only on the encoder and decoder layers that occupy most of the model weights in PipeLoad mechanism design.

For simplicity, we show only three computation ready signals and three memory destruction signals in Fig. 5. And for the sake of clarity, three Loading Agents are symbolized as $LA_{1},LA_{2},LA_{3}$. Model layers are signified as $L_{k}$ where $k$ represents the index within the total number of layers, denoted by $N$. Symbols $S^{\text{comp}}_{k}$, $S^{\text{dest}}_{k}$ and $S^{\text{stop}}$ respectively represent computation ready signal for layer $L_{k}$, memory destruction signal for layer $L_{k}$ and loading stop signal. During the implementation of PipeLoad, the $i$-th Loading Agent is assigned a subset of model layers, following the distribution $L_{i+jm}$, where $i$ ranges from 1 to $m$, with $m$ representing the total number of Loading Agents, and $j$ represents an iterative index, ranging from 0 to $\lfloor(N-i)/m\rfloor$ ($i+jm\leq N$ and $j\in\mathbb{N}$). In this case, $LA_{1}$ is responsible for layers $(L_{1},L_{4},L_{7},\ldots)$, $LA_{2}$ for layers $(L_{2},L_{5},L_{8},\ldots)$ and $LA_{3}$ for layers $(L_{3},L_{6},L_{9},\ldots)$. This layer allocation method is designed to minimize pipeline stalls since we can overlap the inference time of three layers with the loading time of a single layer.

As shown in Fig. 5, the three Loading Agents commence the parallel loading process. As the layer $L_{1}$ is fully loaded to memory, $LA_{1}$ issues the computation ready signal, $S^{\text{comp}}_{1}$ to Inference Agent. After receiving $S^{\text{comp}}_{1}$, Inference Agent starts to perform forward computation for $L_{1}$. If Inference Agent receives $S^{\text{comp}}_{2}$ or $S^{\text{comp}}_{3}$ first, the inference queue in CPU will ensure that the layers are computed in the correct order. Simultaneously, after $LA_{2}$ and $LA_{3}$ load $L_{2}$ and $L_{3}$ as well as sending computation ready signals to Inference Agent, the three Loading Agents are able to continue loading the respective next layers $L_{4},L_{5},L_{6}$ into memory. Following the computation on $L_{1}$, Inference Agent issues the memory destruction signal, $S^{\text{dest}}_{1}$ to Daemon Agent, which initiates the process of de-allocating memory space for $L_{1}$. The loading and inference process for the subsequent layers is similar. In addition, when Daemon Agent detects the system memory usage has exceeded the memory constraints of the edge device, it sends the loading stop signal, $S^{\text{stop}}$ to all Loading Agents, pausing their layers’ loading operation until sufficient memory space is available.

We evaluate the performance and memory footprints of PipeLoad with 2, 4 and 6 Loading Agents and compare them to baseline and to PipeSwitch. We choose these three numbers of Loading Agents since they are essentially factors of the number of encoder or decoder layers in four transformer models. The performance and memory footprints test results are shown in TABLE II and TABLE III respectively, where LAs is an acronym for Loading Agents. In order to show the optimisation results more directly, we add two metrics in tables respectively, the speedup and the ratio, with their expressions are as follows:

We evaluate the performance of Hermes under different memory constraints. In addition, we measure the latency and the corresponding optimal number of Loading Agents.