SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Overhead of Length Prediction Given that we must predict the lengths of all instructions within a group before generating their responses, there is an inherent overhead associated with response length prediction. This overhead entails calculating the keys and values for the instruction tokens and generating a small number of tokens (typically 2 to 3 tokens). However, as depicted in Figure 4, the processing time for instructions is extremely fast, requiring only the time to generate a few tokens
Table 4: Performance of sequence scheduling with different response length perception module.
Throughput (samples/s) ↑	Improvement ↑	Tokens/batch ↓
Vanilla	1.22	377
Ground Truth Predictor	2.52	+107%	201
Pooling + MLP	1.96	+61%	216
[LEN]-token Fine-tune	2.10	+72%	210
Perception Only*	1.40	+15%	328
Instruction Tunning (mean)	1.77	+45%	211
Instruction Tunning (max)	2.27	+86%	208
Table 5: Ablation study on three components of sequence scheduling method.
Binning	FCR	VBS	Throughput (samples/s) ↑	Improvement ↑	Tokens/batch ↓
×	×	×	1.73	+42%	271
×	✓	×	1.76	+44%	222
×	✓	✓	2.22	+82%	209
✓	×	×	1.58	+30%	300
✓	✓	×	1.86	+52%	209
✓	✓	✓	2.27	+86%	203
(ranging from 1 to 4 tokens). Consequently, this overhead can be effectively offset by the overall acceleration achieved through sequence scheduling.
Experimental Setting. Our experiments are conducted on two datasets: a set of 10,000 prompts from a subset of the alpaca dataset [33] (which is different from the one used to train the length predictor) and a set of 429 prompts from the Instruction-in-Wild datasets [36]. The former consists of self-instructed prompts, while the latter contains real-world collected prompts.
For our baseline experiments, we set the batch size to 16. Regarding the variable batch size strategy, we use a batch size of 16 for instructions with a length  $(L)$  greater than or equal to 300. For instructions with a length below 300, we calculate the batch size using the formula  $256 \times 50 / L$  and then round it to the nearest power of 2 for better GPU utilization. We maintain a fixed group size of 256. The inference is performed on the Vicuna-7B [4] model using an 80GB A100 GPU. We sample generations with a temperature of 0.5 for diversity in responses.
We evaluate the throughput at the sample level and report the number of samples processed per second and the corresponding improvement compared to the vanilla baseline, where inference is conducted with batch size 16. The average length refers to the mean length of each batch, where a smaller average length indicates reduced redundant computation.
Results. Table 4 presents the performance of sequence scheduling with different response length perception modules. Among the length predictors considered, the one with instruction tuning demonstrates the best performance, achieving an impressive  $85\%$  improvement in throughput. The acceleration is due to a smaller number of tokens per batch, which includes wasted tokens due to short sentences waiting for the long ones to be completed. It is important to note that the estimation-only method exhibits a significantly higher failure re-collection ratio, resulting in reduced performance. Hence, we report the results obtained using the vanilla sequence scheduling approach.
When predicting the mean length, the failure re-collection ratio increases to  $29.8\%$ , which is considerably higher compared to the  $15.3\%$  achieved when predicting the maximum length. Consequently, the performance improvement drops to  $45\%$ . Alternatively, if we utilize the ground truth (i.e., the maximum length observed during multiple inferences) as the length predictor, it serves as an upper bound for this method. Our approach performs only 0.25 samples/s slower than the upper bound, showcasing that an effective length predictor can yield substantial improvements.
Table 6: Performance of sequence scheduling on Instruction-in-Wild dataset.
Throughput (samples/s) ↑	Avg. length ↓	Error ↓	Acc-50 ↑	Acc-100 ↑
Vanilla	0.78	475	-	-	-
Estimation Only	1.07	422	358	20%	38%
Instruction Tunning	1.24	299	139	43%	68%
Table 7: Study on the Failure Collection and Recomputation (FCR) ratio.
k	Avg. Wait	Max. Wait	FCR ratio	FCR time	Throughput (samples/s)
-50	42%	63%	34%	21%	1.80
-10	37%	53%	20%	20%	2.10
0	37%	52%	15%	20%	2.27
+10	45%	52%	14%	14%	2.20
+50	38%	55%	7%	7%	2.10
Furthermore, our method exhibits a low variance in throughput, with a value of 0.05 over three times experiments. This indicates the stability and consistency of our approach in achieving improved inference speed.
Table 5 presents an ablation study on the three components of our approach, demonstrating their individual contributions to the overall improvement. We observe that Binning enhances the effectiveness of Variable Batch Size (VBS), leading to a more powerful combination. Each component plays a significant role in achieving the final improvement in throughput. Furthermore, Table 6 showcases the performance of our methods on the Instruction-in-Wild dataset, confirming the effectiveness of our approach across different datasets.
In addition, on the right side of Figure 3b, we analyze the relationship between throughput and group size. We observe that a larger group size provides more flexibility for sequence scheduling, resulting in improved throughput. However, the rate of improvement gradually slows down as the group size increases. Thus, it becomes crucial to strike a balance between the overhead involved in gathering the group size information in real-world scenarios and the achieved improvement.
Waiting time. Apart from the throughput, another important metric is the waiting time for each user. We introduce three metrics for comparison: maximum waiting time, average waiting time, and their ratio (maximum waiting time multiplier). The maximum timing time is a fairness metric for measuring individual request delay. We define the waiting time for a user to be the delay from receiving the request (start of processing a group) to generate the corresponding response. With group size 256, our method saves  $63\%$  average waiting time compared to the vanilla. Due to the FCR mechanism, our maximum waiting time multiplier is 2.7, which is 1.4 times the vanilla's one. However, with inference speed acceleration, the maximum wait time is also reduced by  $48\%$ .
FCR ratio. One factor influencing inference latency is the Failure Collection and Recomputation (FCR) ratio. This ratio represents the proportion of FCR samples recalculated at the end of a batch, which causes delays in processing. To assess the effect of different FCR ratios, we modify the predicted length by a constant value  $k$ . A larger predicted length results in a lower FCR ratio (more tolerable). However, a long response with a short predicted length may introduce more waste generation in batch (generating  $k$  more times). The experimental results are as shown in Table 7. The Avg. Wait and Max. Wait columns denote the the average and maximum wait times compared to the vanilla method. FCR time indicates the proportion of time utilized for FCR processing.
When using an accurate length predictor, it is advisable to directly utilize the predicted length with a  $15\%$  FCR ratio. This approach strikes a balance between the time spent on FCR recomputation and the time wasted during batch computation. In some cases, the generation process might require just a few more tokens to complete the task, but it is difficult to determine which response will finish or if it will finish even after generating more tokens. This concern is also the reason why the maximum length among four times generation is used, as it aims to reduce the potential FCR ratio.
5 Limitation and Discussion
One limitation is that accurate response length prediction is not always guaranteed, even with the instruction-tuned length predictor module we developed. While our experiments demonstrate promising results, there is still room for improvement in the accuracy of response length estimation.
Besides, although sequence scheduling significantly reduces redundant computations, it cannot completely eliminate the issue. Even with a ground truth response length predictor, the ratio of redundancy decreases from  $66\%$  to  $33\%$ , leaving ample room for further improvement. Recent works such as ORCA [37] have proposed novel inference systems that offer alternative solutions to mitigate the redundant computation problem.
Another limitation of our work is that we focus less on the process of input instructions. As the maximum token length supported by LLMs increases, users may input very long instructions, leading to redundant computations and inefficiencies. Future work could explore a combination of our proposed sequence scheduling with input batch scheduling techniques, such as those discussed in Fang et al. (2021) [10], to further optimize the inference process.
Our approach assumes that the group size exceeds the capacity of a single GPU. As large language models may become a ubiquitous infrastructure, similar to search engines, the number of queries will increase significantly. Furthermore, the emergence of models like GPT-4 with 32k sequence length support and Claude with 100K sequence length support amplifies the challenge of handling varying response lengths, highlighting the relevance and effectiveness of our method.
Our approach is easily extensible to multi-GPU settings, where multiple GPUs are used for faster processing and handling larger model sizes. By reallocating batches of sequences with different perceived response lengths to different GPU nodes, our method remains effective in high-performance computing environments. This extension ensures scalability and maintains the efficiency gains achieved in our approach.
Some argue our method needs to reorder inputs. However, the order of input sequences is not critical; rather, the focus is on batch assembly. If requests in a group arrive together, the overhead from reordering indices is negligible. The main latency arises from the response length perception, not index reordering. After forming batches, their processing order is flexible. Although this might affect individual user latency, it's bound by a group's processing time. Batch shuffling can also offset order impacts. In multi-GPU setups, batches with varied predicted response lengths are sent to different GPUs for simultaneous processing, emphasizing batch assembly over sequence order.
For more powerful LLMs (such as GPT-4, Claude, etc.), our proposed PiA method introduces no need for instruction-tunning and introduces no overhead during sequence-scheduling. This can better handle long inputs and multi-turn conversation. However, as these models are not open-sourced, it is unable to perform experiments on them. Detailed discussion on PiA method can be found in appendix C.
Our findings indicate that large language models (LLMs) possess a comprehensive comprehension of their generated responses. This understanding opens up possibilities for the creation of faster inference techniques, including non-autoregressive methods, that can overcome the performance constraints associated with sequential token generation.
In this paper, we proposed a novel technique called sequence scheduling, which optimizes large language model (LLM) inference by leveraging response length prediction. Our approach groups queries with similar response lengths in batches, reducing computational waste and improving inference throughput. We introduced the concepts of failure re-collection and variable batch size to further enhance efficiency. Experimental results on real-world instruction datasets using the Vicuna-7B model demonstrated an  $86\%$  improvement in throughput without sacrificing performance quality. Our method offers a valuable addition to the existing toolkit for LLM inference, addressing the challenge of efficient deployment of LLMs at scale.
Yang You's research group is being sponsored by NUS startup grant (Presidential Young Professorship), Singapore MOE Tier-1 grant, ByteDance grant, ARCTIC grant, SMI grant and Alibaba grant. This work is sponsored by Huawei Noah's Ark Lab. We would like to thank Shenggan Cheng and Ziming Liu for their valuable discussion, including advice on the implementation.
[22] S. Li et al., “Colossal-ai: A unified deep learning system for large-scale parallel training,” arXiv preprint arXiv:2110.14883, 2021.
[23] C. Lyu, J. Xu, and L. Wang, "New trends in machine translation using large language models: Case examples with chatgpt," arXiv preprint arXiv:2305.01181, 2023.
[24] X. Ma, C. Zhou, X. Li, G. Neubig, and E. Hovy, "Flowseq: Non-autoregressive conditional sequence generation with generative flow," arXiv preprint arXiv:1909.02480, 2019.
[25] L. Ouyang et al., "Training language models to follow instructions with human feedback," Advances in Neural Information Processing Systems, vol. 35, pp. 27 730-27 744, 2022.
[26] A. Paszke et al., "Automatic differentiation in pytorch," 2017.
[27] R. Pope et al., "Efficiently scaling transformer inference," arXiv preprint arXiv:2211.05102, 2022.
[28] J. Robinson, C. M. Ryting, and D. Wingate, “Leveraging large language models for multiple choice question answering,” arXiv preprint arXiv:2210.12353, 2022.
[29] Y. Sheng et al., "High-throughput generative inference of large language models with a singlegpu," arXiv preprint arXiv:2303.06865, 2023.
[30] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” arXiv preprint arXiv:1909.08053, 2019.
[31] R. Shu, J. Lee, H. Nakayama, and K. Cho, “Latent-variable non-autoregressive neural machine translation with deterministic inference using a delta posterior,” in Proceedings of the aaai conference on artificial intelligence, vol. 34, 2020, pp. 8846–8853.
[32] Z. Sun, Z. Li, H. Wang, D. He, Z. Lin, and Z. Deng, "Fast structured decoding for sequence models," Advances in Neural Information Processing Systems, vol. 32, 2019.
[33] R. Taori et al., Stanford alpaca: An instruction-following llama model, https://github.com/tatsu-lab/stanford_alpaca, 2023.
[34] H. Touvron et al., "Llama: Open and efficient foundation language models," arXiv preprint arXiv:2302.13971, 2023.
[35] X. Wu, C. Li, R. Y. Aminabadi, Z. Yao, and Y. He, "Understanding int4 quantization for transformer models: Latency speedup, composability, and failure cases," arXiv preprint arXiv:2301.12017, 2023.
[36] F. Xue, K. Jain, M. H. Shah, Z. Zheng, and Y. You, Instruction in the wild: A user-based instruction dataset, https://github.com/XueFuzhao/InstructionWild, 2023.
[37] G.-I. Yu, J. S. Jeong, G.-W. Kim, S. Kim, and B.-G. Chun, “Orca: A distributed serving system for {transformer-based} generative models,” in 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), 2022, pp. 521–538.
[38] T. Zhang, F. Ladhak, E. Durmus, P. Liang, K. McKeown, and T. B. Hashimoto, "Benchmarking large language models for news summarization," arXiv preprint arXiv:2301.13848, 2023.
A Additional Instruction Tuning Implementation Details
To fine-tune the instruction-tuned length predictor module, we followed a specific procedure. First, we prepared the dataset by sampling each instruction four times. One sample was obtained using greedy decoding with a temperature of 0, while the other three samples were generated using a temperature of 1.0 with different random seeds. The maximum length among the four samples was used as the target length for each instruction.
For the training dataset, we constructed new instructions by appending the requirement of perceiving the response length only. The prompt we used for this purpose was: "Don't output the response for the above instruction. Instead, you need to predict the number of tokens in your response. Output only one number."
During the training process, we employed the same hyperparameters as the Vicuna [4] instruction tuning process on LLaMA [34]. Specifically, we set the learning rate to 0.00005 and trained the model for three epochs. We applied the LoRA [16] method solely to the query and key linear layer. The training was conducted on a single 80GB A100 GPU. All codes are implemented in PyTorch [26].
B Distribution of Instruction-in-Wild Dataset
The histogram of response length on Instruction-in-Wild [36] dataset is shown in Figure 5. The response length is more diverse and contains more responses with a long length.
Figure 5: Response length distribution of Instruction-in-Wild dataset.
C Discussion on Sequence Scheduling with PiA
In the main text, we presented the sequence scheduling technique using instruction-tuned models, where the LoRA weight was utilized for response length perception. However, recent LLMs such as GPT-4 and Claude have shown promising performance in Perception in Advance (PiA), which allows them to leverage PiA for sequence scheduling without the need for additional weights.
To understand the PiA process, consider the following format:
[Input] [PiA Prompt] [PiA Response] [Response]
Here, the PiA prompt might be something like "Please estimate the length of your prediction," and the subsequent PiA Response could be "The estimated response length is [X]". The primary overhead of this method involves the generation of the aforementioned tokens, typically no more than 30 in number.
As stated section 3.1, PiA may have different behaviour than Perception Only (PO) because the response is now depend on PiA prompt and response. However, we can eliminate this effect due to the fact that the input and the response can be detached from the PiA Prompt and its Response.
To ensure the quality of the response remains unchanged, we can use the following two methods:
- After the generation of the PiA response, we can exclude both the PiA Prompt and PiA Response embeddings. The generation of the primary Response would then only use the kv-cache of the input.
- Alternatively, the Response can be generated with both the PiA Prompt and PiA Response masked.
Both these strategies ensure the preservation of generation quality. In our experiments with models like ChatGPT and other leading-edge Large Language Models (LLMs), merely invoking the API didn't allow for the bypassing of the PiA prompt's influence. However, this implementation is quite simple for the companies. This observation also led us to introduce the Perception Only (PO) approach, wherein we compared the mean predicted length against the original prediction length as an estimation of performance of PiA method in practice.
To further improve the speed of inference in this setting, one potential approach is to reuse the key-value (kv) cache of input queries from the response length perception stage during the inference stage. This eliminates the need for recomputing the kv-cache on input instructions, thereby saving valuable processing time.
One strategy we explored is offloading the kv-cache to the main memory and then loading it back into the GPU memory. However, the transfer time between the CPU and GPU can be significant and greatly impact overall performance, often surpassing the time saved by recomputation. To address this, one possible improvement is to offload and load only a portion of the kv-cache asynchronously, reducing the transfer time overhead. This is an area that we leave for future work and exploration.
Another approach we investigated involved compressing the kv-cache and storing it directly in the GPU memory. We applied FlexGen's quantization method [29] for compression and found that it had minimal impact on performance. However, this approach does consume additional memory and can lead to a smaller batch size, potentially degrading overall performance. A potential avenue for further exploration is to combine compression and offloading techniques to strike a balance between memory usage and performance.
Considering these factors, we have chosen to continue using the recomputation strategy for the PiA response length perception module in our proposed pipeline. While there is potential for optimizations through offloading and compression, further investigation and refinement are required to achieve substantial performance gains in practice.