SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills

解决问题：该论文旨在解决大语言模型（LLM）推理过程中的低计算利用率和微批次之间的不平衡，提出了一种名为SARATHI的方法。

关键思路：SARATHI采用了分块预填充和解码最大批处理的方法。分块预填充将预填充请求分成相同大小的块，解码最大批处理使用单个预填充块构建一个批次，并使用解码填充其余插槽。在推理过程中，预填充块饱和GPU计算，而解码请求则“附加”在上面，与仅解码批次相比成本低得多。分块预填充允许从单个预填充请求构建多个解码最大批次，最大化可以附加的解码覆盖范围。此外，这些批次的均匀计算设计减轻了微批次之间的不平衡，显著减少了管道气泡。

其他亮点：该论文的实验结果表明，SARATHI在不同模型和硬件上都取得了显著的推理性能提升。例如，在A6000 GPU上使用LLaMA-13B模型，SARATHI将解码吞吐量提高了多达10倍，并将端到端吞吐量加速了多达1.33倍。在A100 GPU上使用LLaMa-33B模型，我们实现了高达1.25倍的端到端吞吐量和多达4.25倍的解码吞吐量提高。当与管道并行处理一起使用时，SARATHI将气泡减少了6.29倍，从而将端到端吞吐量提高了1.91倍。

相关研究：最近的相关研究包括：

  • "Scaling Laws for Neural Language Models" by Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, Dario Amodei, and OpenAI (2020)
  • "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism" by Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro, and Andrew Gibiansky (2019)
  • "GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding" by Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, and Christopher A. Cho (2020)

论文摘要：SARATHI: 通过与分块预填充一起使用的解码来实现高效的LLM推理 作者：Amey Agrawal，Ashish Panwar，Jayashree Mohan，Nipun Kwatra，Bhargav S. Gulavani，Ramachandran Ramjee 大型语言模型（LLM）推理由两个不同的阶段组成 - 预填充阶段处理输入提示，解码阶段自回归地生成输出令牌。虽然预填充阶段在小批量大小下有效地饱和了GPU计算，但解码阶段每次请求只生成一个令牌，导致计算利用率低。同时，预填充和解码时间的变化也会导致使用管道并行性时微批次之间的不平衡，进而导致气泡等效率低下。 我们提出了SARATHI来解决这些挑战。SARATHI采用分块预填充，将预填充请求分成相等大小的块，并使用解码最大批处理构建批处理，该批处理使用单个预填充块并填充其余插槽。在推理期间，预填充块饱和了GPU计算，而解码请求“搭便车”并且与仅解码批处理相比成本低至一个数量级。分块预填充允许从单个预填充请求构建多个解码最大批处理，最大化可以搭便车的解码覆盖范围。此外，这些批处理的统一计算设计改善了微批次之间的不平衡，显著减少了管道气泡。 我们的技术在各种模型和硬件上显著提高了推理性能。对于A6000 GPU上的LLaMA-13B模型，SARATHI将解码吞吐量提高了高达10倍，并将端到端吞吐量加速了高达1.33倍。对于A100 GPU上的LLaMa-33B，我们实现了1.25倍的端到端吞吐量和高达4.25倍的解码吞吐量。在GPT-3上与管道并行性一起使用时，SARATHI将气泡减少了6.29倍，从而使端到端吞吐量提高了1.91倍。

内容中包含的图片若涉及版权问题，请及时与我们联系删除