如何进行高效的LLM分布式推理

1.背景

这次指的大型语言模型（LLMs）主要说的是采用Transformers架构的模型，该架构在训练阶段具有高度的并行性。然而，在实际应用中，高效地部署这些模型面临一些挑战。这是因为生成推理是逐个token进行的，而每个token的计算都依赖于先前生成的标记。因此，支持大规模训练的模型需要仔细考虑并行布局和内存优化，以实现所需的可扩展性和高效的低延迟推理。从而更好地支持大规模训练的模型在实际应用中实现高效的低延迟推理，从而满足对于实时性和可扩展性的需求。

然而，在实际的工程化环境中，模型性能优化从未是一个简单的最优化问题。在大型语言模型（LLMs）的生成推理中，我们面临着诸多挑战，包括庞大的内存占用、严格的延迟目标以及长序列长度，不同指标质检的优化方法还存在一定的冲突。这时就需要引入了帕累托前沿的概念，用以描述一组解决方案，即在不牺牲其他目标的前提下，无法改善一个特定目标。

• 大模型训练/推理的主要特点：large deep models（模型大）, with tight latency targets（时间紧） 以及 long sequence lengths（理解长度长）；
• Pareto frontier：核心思想是，在多目标决策问题中，存在一组解决方案，这些解决方案在某些目标上的改善是不可能的，而在其他目标上的改善是可能的。这意味着你不能在一个目标上取得改进而不牺牲其他目标。这些解决方案构成了 Pareto frontier，也称为 Pareto 边界。在 Pareto frontier 上的解决方案通常被认为是非劣解（non-dominated solutions）或 Pareto 最优解（Pareto optimal solutions）。举例来说，考虑一个制造公司需要在成本和质量两个目标上做决策。如果提高产品质量会增加成本，那么在成本和质量之间存在一个权衡。Pareto frontier 就是描述了在不牺牲一方面目标的情况下，如何最大化或最小化另一目标。

参考《Efficiently Scaling Transformer Inference》提供的一套工程原则，可以引导Transformer结构的推理并行策略优化，以帮助我们在考虑模型规模和特定应用需求的情况下，经验性地推导出最佳的模型并行策略。对于不同的推理需求，可以通过实验列表进行按需调整，选择合适的核心数和批处理大小，如下图，可以根据在不同场景中的延迟要求，选择最优的核心数（加速卡数量）以及批处理大小：

▲ 图 1：《Efficiently Scaling Transformer Inference》

2.推理成本权衡

一次端到端的推理延迟需要考虑到两个主要方面：数据处理（prefill，大部分在CPU上进行处理）和逐词生成（decode，大部分在GPU上进行）。

衡量指标：为了衡量吞吐量，我们可以使用MFU（The Model FLOPS Utilization）作为指标，即实际吞吐量与理论最大吞吐量的比值（也可以用于衡量计算效率）。

内存成本：在prefill和decoder两个阶段都会存在，主要体现在将数据从HBM（high-bandwidth memory）传输到计算核心的成本上。

计算成本：一般涉及到对于一个包含N个参数的decoder-only模型来说，生成一个token需要进行2N次矩阵乘法运算。计算成本即指这些乘法运算的代价。

3.分布式推理

1）符号表示

（1）  ：表示最后一个维度E已经被切分为X * Y * Z个分区，其中 BLE分别表示：Batch、Sequence Length、Model Embed；此外F表示MLP前向传播中的维度。

（2）后缀“partialsum-x”：表示一个指定的Tensor在每个Chip上被本地合并（contracted/summed）了，但仍需要在TPU x轴上对芯片进行求和才是最终结果。

（3）all-reduce(x)：表示在x这个维度上对  这样的切片进行合并，最终输出  ，all-reduce = reduce-scatter + all-gather；

（4）reduce-scatter(x)：对x维度上芯片的张量进行求和，并再该维度上对求和后的张量进行重新分片，相当于消除掉partialsum-x并在任意维度增加一个下标x；

（5）all-gather(x)：可以直接理解为concat，可以消除掉某一个维度的下标；

（6）all-to-all(x)：相当于转置，可以实现  ->  这样的操作；

▲ 图 2：常见的分布式通讯方法 

2）FFN的并行

随着在更多芯片上并行计算，内存延迟和计算延迟通常会呈近线性下降。然而，通信延迟依然会保持大致不变，因为对于每一对矩阵乘法运算来说，整个激活矩阵是跨芯片聚合的。因此通讯延迟会成为超大规模分布式训练的瓶颈；

其中FFN的并行方法大致可以分为三类：

▲ 图 3：FFN分布式并行方法

▲ 图 4：不同batch下不同并行策略的消耗

3.注意力层的并行

目前注意力的并行机制主要包括了MHA以及MQA（https://arxiv.org/pdf/1911.02150.pdf，19年Google提出的一种新的Attention机制，可以加快decoder的生成）的两种Attention Layer的并行，其主要并行方式如下：

▲ 图 5：Attention 并行方式

▲ 图 6：MHA&MQA并行表示

MHA的切分和FFN Partion相似，不同之处在于将F维度替换为H维度进行切分。如图6(a)所示，可以选择仅切分H维度。在处理器数不能整除注意力头数的情况下，特殊处理是必要的。对于处理器数量大于注意力头数的情况，注意力头的部分会进行部分复制（partially replicated），与FFN类似，这等效于一次AllReduce=ReduceScatter+Allgather。

MQA的切分相对较为复杂。如果仍然选择在H维度上进行切分，如图6(b)方案所示，K和V张量在所有注意力头之间是共享的，但它们必须在每个芯片上进行复制，这会导致MQA的内存成本节省失效。因此，作者采用图5(c)切分，即对B维度进行切分。然而，这也要求任务能够组织成较大的batch，以便有足够的切分余地。具体而言，Q、K和V矩阵在批次B维度上被分成N个分区，这降低了每个芯片加载KV缓存的内存成本，也因此减少了相应的访存。与MHA的并行策略相比，MQA需要使用all-to-all对进行输入输出激活张量resharding，从而带来额外的通信成本。值得注意的是，使用较大的Batch Size时，MQA可以减少内存访问成本，但同时会增加并行通信开销。

总的来说，MQA允许使用更大的批处理大小和序列长度，从而在减少内存时间的同时降低延迟，提高吞吐量（训练更快）。

4.其他优化策略

1） 使用PaLM中并行的FFN以及Attention层：可以将两层layerNorm优化成一层，并且对于FFN层的输入矩阵、KV的投影矩阵、输出矩阵都可以实现复用，从而提高模型整体的吞吐率；

2）Low-level优化：主要指的是使用了Looped CollectiveEinsum（https://dl.acm.org/doi/abs/10.1145/3567955.3567959）的技术，可以降低通讯成本；

4.总结

本文系统性地阐述了稠密LLM模型并行推理的目标、方法和规律。针对稠密LLM模型（主要指的是Transformer类的模型）并行推理的目标、方法和规律该论文进行详细的总结。

首先，在Prefill和Decoding阶段的并行策略，在模型推理过程中具有不同的特性和需求。Prefill主要发生在CPU上，而Decoding则主要在GPU上进行处理。因此，针对不同的计算资源和优化需求，对这两个阶段的并行化需要采用不同的策略。

其次，强调了在Low Latency（低延迟）和High Throughput（高吞吐量）的场景下，并行策略也呈现出不同的特征。在追求低延迟的情境下，需要特别关注推理速度和实时性。而在高吞吐量的情境下，更侧重于系统整体性能的提升。因此，针对不同的应用场景，需要灵活选择并实施不同的并行化策略以满足特定需求。

最后，对于MQA（Mixed Quantization Attention）在并行Attention和Feed-Forward Network（FFN）方面的优势进行了重点强调。MQA的采用对于推理性能提升具有显著裨益，特别是在处理注意力机制和前馈网络时。这进一步强调了在不同任务和模型结构下，采用差异化的并行化策略对于性能优化至关重要。

综合而言，稠密LLM模型的高效并行推理具备一定的复杂性，不同阶段、不同场景下的并行策略有不同的指导原则。MQA的引入更进一步丰富了并行计算的工具箱，为提升模型推理效率提供了有力支持。

参考文献

[1] Pope R, Douglas S, Chowdhery A, et al. Efficiently scaling transformer inference[J]. Proceedings of Machine Learning and Systems, 2023, 5.

[2]​ https://zhuanlan.zhihu.com/p/634236135

[3]​ https://www.zhihu.com/question/591646269/answer/3246726206

[4] Wang S, Wei J, Sabne A, et al. Overlap communication with dependent computation via decomposition in large deep learning models[C]//Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. 2022: 93-106.

[5] Shazeer N. Fast transformer decoding: One write-head is all you need[J]. arXiv preprint arXiv:1911.02150, 2019.

本文转载自 ​​AI遇见云​​，作者： 杨希