HexGen: Generative Inference of Large Language Model
over Heterogeneous Environment

Notations. We first introduce the following notations: Let $\mathbf{D}=\{d_{1}\ldots d_{N}\}$ be a set of $N$ GPU devices, where $M_{d}$ denotes GPU memory limit, $m_{d}$ denotes GPU memory bandwidth, and $c_{d}$ denotes tensor core computation power; $\mathbf{A}\in\mathbb{R}_{+}^{N\times N}$ and $\mathbf{B}\in\mathbb{R}_{+}^{N\times N}$ be the communication matrix between these devices describing the latency and bandwidth respectively, where the latency and bandwidth between device $d$ and $d^{\prime}$ is $\alpha_{d,d^{\prime}}$ and $\beta_{d,d^{\prime}}$; $L$ denotes the total number of layers in the model to be served. We summarize all notations used in this paper in Appendix A for easy reference.

Formalization of the scheduling problem. Given the above notations, we can formalize our scheduling problem as follows: suppose $\mathbf{d}_{i,j}$ is a subset of GPU devices that satisfies $\bigcup_{i,j}\mathbf{d}_{i,j}\subseteq\mathbf{D}$, and $\mathbf{d}_{i,j}\bigcap\mathbf{d}_{i^{\prime},j^{\prime}}=\emptyset,\quad\forall i\neq i^{\prime}\vee j\neq j^{\prime}$, where we suppose that the union of subsets of GPUs $\bigcup_{j}\mathbf{d}_{i,j}$ serves the $i$-th model replica as an independent pipeline, and the set of GPUs $\mathbf{d}_{i,j}$ serves the $j$-th stage in the $i$-th pipeline that holds $l_{i,j}$ transformer layers, if $\|\mathbf{d}_{i,j}\|>1$, we indicate that the subset of GPUs $\mathbf{d}_{i,j}$ runs tensor model parallelism. An assignment $\sigma$ is a mapping: $\mathbf{D}\rightarrow\left\{\left(\mathbf{d}_{i,j},l_{i,j}\right)\right\}$, corresponding to a layout of multiple inference worker groups to serve multiple replicas of the same model simultaneously. Consider a set of inference tasks $\mathbf{T}$ that satisfy some distribution¹¹1The most commonly-used case is Poisson distribution; while one can switch it for any particular distributions. $\mathbf{T}\sim\mathcal{P}$. An optimal assignment $\sigma^{*}$ can be defined as:

where $\textsc{C}_{\text{comm}}\left(\sigma\right)$, $\textsc{C}_{\text{comp}}\left(\sigma\right)$, $\textsc{C}_{\text{mem}}^{d}\left(\sigma\right)$ represents communication cost, computation cost, and memory limit for layout $\sigma$. Intuitively, the scheduling problem is to find an optimal assignment that partitions the device set to represent multiple independent inference pipeline groups that can maximize the inference service level objective (SLO) considering the computation cost, communication cost, and memory consumption constraints. We summarize the estimation of the computation cost, communication cost, and memory consumption constraints in Table 1, and the detailed formulation is enumerated in Appendix B. As one can see, solving this problem is obviously NP-hard; thus, we adopt a two-phase search algorithm to tackle the problem:

Next, we introduce the following dynamic programming (DP) algorithm to solve this problem. We define a memorization buffer $\textsc{DP}\in\mathcal{R}^{S_{i}\times\#_{1}\times\#_{2}...\times\#_{N_{T}}}$, where $S_{i}$ is the total number of stages in this pipeline, $\#_{k},k=1,...,N_{T}$ is the total number of $k$-th type of GPUs in the GPU set $\mathbf{d}_{i,\sim}$, the value $\textsc{DP}\left[j;\boldsymbol{\tau}\right]$ represents the communication and computation cost of assigning the first $j$ stage(s) in the GPU set represented by $\boldsymbol{\tau}$ if the memory limit for the first $j$ stage(s) is satisfied²²2$\textsc{DP}\left[j;\boldsymbol{\tau}\right]=+\infty$ if the memory limit is violated for any stage.. For example, $\textsc{DP}\left[1;\left[2,0,...,0\right]\right]$ represents the communication and computation cost of assigning the first stage in two GPUs of the first type. If we initialize the values in DP to $+\infty$ and set $\textsc{DP}\left[0;\boldsymbol{\tau}\right]=0$, we can derive the transition formula as:

where $\tau_{k}\cdot\mathbf{e}_{k}$ is the vector representation of a set of $\tau_{k}$ $k$-th type GPUs, $\textsc{C}^{i,j}_{\text{comp}}\left(\tau_{k}\cdot\mathbf{e}_{k}\right)$ is defined by Table 1, and $\textsc{C}^{i,j}_{\text{comm}}\left(\tau_{k}\cdot\mathbf{e}_{k}\right)$ is the sum of tensor parallel communication cost and pipeline parallel communication cost follow the formulation in Table 1. Notice these two value will be set to $+\infty$ if the memory limit $\textsc{C}_{\text{mem}}^{d}\left(\tau_{k}\cdot\mathbf{e}_{k}\right)$ defined in Table 1 is violated. Formally, we formalize a recursive implementation of the dynamic programming algorithm in Algorithm 1 to estimate the minimal cost of this pipeline — the assignment of the pipeline stage can be implemented by a standard back-track process over the memorization buffer DP straightforwardly, which we do not enumerate here.

Complexity analysis. The computation of the cost defined in Table 1 can be done in constant time. For each stage, the proposed DP algorithm will visit at most $\prod_{k=1}^{N_{T}}{\#_{k}}$ different subsets of GPUs; the maximal depth of the recursion is $S_{i}$; thus, the total time complexity of the proposed DP algorithm is $\mathbb{O}\left(S_{i}\cdot\prod_{k=1}^{N_{T}}{\#_{k}}\right)$. This is much more efficient than the vanilla approach based on enumeration without the heuristic of forcing tensor parallelism to use the same type of GPUs — in that case, each stage has to consider $2^{\left|\mathbf{d}_{i,~{}\sim}\right|}$ different subsets of GPUs, where the total time complexity grows to $\mathbb{O}\left(S_{i}\cdot 2^{\left|\mathbf{d}_{i,~{}\sim}\right|}\right)$. In practice, even a highly heterogeneous GPU pool probably only includes a limited variety of GPU types. The proposed DP algorithm can solve the sub-optimization problem efficiently. This process can be further accelerated by limiting the degree of tensor model parallelism to a smaller candidate set, e.g. $\{1,2,4,8\}$, which reduces the total time complexity to $\mathbb{O}\left(S_{i}\cdot 4^{N_{T}}\right)$.

Initialization. A good initialization of the genetic algorithm is usually helpful in accelerating the whole search process. Thus, we initialize the population with a simple heuristic based on the communication condition. Formally, given the communication matrix $\mathbf{A}$ and $\mathbf{B}$, we execute the vanilla K-means algorithm to construct $M$ independent pipeline parallel groups, where the hyper-parameter $M$ is determined by the standard Elbow method. Intuitively, this helps the assignment avoid using slow cross-region communication links. Note that $M$ is not fixed after initialization, as we can change this by the mutation operations introduced next.

Mutation operations. We follow the vector representation of a set of GPUs introduced earlier in Section 4.2 and define three mutation operations in the genetic algorithm.

Determine the pipeline partitions. We adopt a simple expectation maximization algorithm to determine any particular pipeline stage partition (i.e. $\left\{l_{i,1},...,l_{i,j},l_{i,S_{i}}\right\}$). Concretely, we first generate an even partition for any new off-spring generated by the mutation operations ($l_{i,j}=\frac{L}{S_{i}},j=1,2,...,S_{i}$). Then after running the dynamic programming algorithm (Algorithm 1), we adjust the pipeline partition proportional to the total memory of the GPU set that currently serves this stage. We find that this heuristic is effective to find a good pipeline partition.

Put it together. We adopt a standard iterative procedure of the genetic algorithm: for each iteration, we conduct the mutation operation(s), run the dynamic algorithm to determine the optimal pipeline assignments, and adjust the pipeline partitions for the current off-springs. To estimate the expected SLO, we adopt the inference task simulator from AlpaServe (Li et al., 2023).

Baselines. We carefully select state-of-the-art approaches as baselines. To understand the cost performance trade-offs and system efficiency of HexGen, we compare: i) HexGen under heterogeneous-full-price, ii) a truncated version of HexGen without asymmetric parallel support under heterogeneous-full-price (the allocation of model replicas is still scheduled by our proposed search algorithm), iii) HexGen under heterogeneous-half-price, and iv) FlashAttention (Dao, 2023) under homogeneous-data-center setting.⁴⁴4Our implementation is identical to the standard FlashAttention implementation under a homogeneous setting. To understand end-to-end performance, we compare HexGen with Huggingface-TGI (HuggingFace, 2023) as the state-of-the-art approach under the homogeneous setting and Petals (Borzunov et al., 2023) as the state-of-the-art approach under decentralized heterogeneous setting. To understand the efficiency of the proposed scheduling algorithm, we compare its convergence with a strawman policy based on random mutation.

Evaluation metrics. Following the generative inference evaluation setup from AlpaServe (Li et al., 2023), we test system performance based on SLO attainment, refers to the proportion of requests that can be finished within a predefined performance threshold. Specifically, we measure SLO attainment in terms of the percentage of requests served within the time frame set by the SLO. We scale the SLO to various multiples of the execution latency of A100 GPUs (SLO Scale in Figure 2), which allows us to evaluate system performance under different levels of operational stringency. We generate inference workload according to a Poisson process parameterized by the request rate. The consecutive requests (inter-arrival times) follow an exponential distribution. For a target SLO goal (e.g., $99\%$), we focus on two metrics: i) the minimum latency deadline required to achieve the desired attainment, and ii) the system’s resilience to peak request rate. We apply the most popular open-source Llama-2 (70B) model on some real-world prompts (Lmsys, 2023), and test output sequence lengths from $32$ to $128$, and request rates varying between $0.125$ - $10$ requests per second.

Cost efficiency. When given the relatively same budget, HexGen under the full budget in the heterogeneous setting clearly outperforms FlashAttention in the homogeneous datacenter setting. In fact, HexGen reaches up to $2.3\times$ and on average $1.5\times$ lower latency deadlines, and is capable of handling a peak request rate that is up to $4\times$ higher ($2\times$ higher on average). Specifically, by analyzing the scheduling results⁵⁵5All the strategies chosen by the scheduling algorithm can be found in Appendix F. for heterogeneous-full-price cases, we find that our scheduling approach always prioritizes intra-machine tensor model parallelism to minimize single request latency and employs inter-machine pipeline parallelism to reduce communication over limited bandwidth. It avoids cross-region communication due to ultra-low bandwidth and aims to maximize device memory utilization by incorporating as many model replicas as possible, thereby enhancing parallel request processing. Furthermore, even when we reduce the budget in the heterogeneous setting by half, HexGen still reveals similar performance to FlashAttention in the homogeneous setting. We believe that this is strong evidence to illustrate that a decentralized system such as HexGen is capable of managing heterogeneous GPUs to provide more economical foundation model inference services without compromising service quality.

Asymmetric parallelism implementation. We also conduct a group of benchmarks to compare HexGen w/wo asymmetric parallel group support under the full budget in the heterogeneous setting. The experimental results reveal that the asymmetric parallelism implementation results in up to $1.8\times$ improvement in terms of reaching lower latency deadlines than the original symmetric implementation from FlashAttention. Additionally, the asymmetric parallelism implementation can manage up to $2\times$ higher peak traffic request rate ($1.5\times$ on average) compared to the symmetric counterpart. This indicates that besides effective scheduling, the asymmetric parallelism implementation is also necessary to unleash the potential of heterogeneous computational power.

Compare with Petals. We first compare HexGen with Petals, the state-of-the-art decentralized inference service engine. Figure 3 illustrates the comparison — under the half-price budget in the heterogeneous setting, HexGen outperforms Petals significantly by achieving up to $3.5\times$ lower latency deadline and managing to handle requests at up to $10\times$ higher rates. This is strong evidence to justify the design of HexGen: Petals is mainly built on top of swarm parallelism (Ryabinin et al., 2023), which heavily depends on dynamic adjustment of the collective learning paradigm to ensure the elasticity in a decentralized machine learning system; however, such a dynamic design compromises the inference service performance significantly when comparing with a system like HexGen that is equipped with the careful design of static scheduling of the inference workflow.

Compare with Huggingface-TGI. We also compare HexGen under the full budget in heterogeneous setting with Huggingface-TGI in the homogeneous data center settings. In this case, HexGen gets almost the same end-to-end performance in both latency and request handling, achieves up to $1.25\times$ lower latency deadlines, and is able to handle requests at the same level of rates.