分类： 未分类

近年来，巨型语言模型（LLM）在自然语言处理领域取得了显著进展，其参数规模已突破千亿甚至万亿大关。然而，这些模型的推理效率却面临着巨大的挑战，尤其是当需要处理长序列文本并满足严格的延迟要求时。本文将深入探讨如何通过模型分割和优化策略来提高巨型语言模型的推理效率。

推理效率的挑战

巨型语言模型的推理效率主要受以下因素影响：

• 庞大的内存占用： 训练好的模型参数和解码过程中的中间状态都需要占用大量的内存空间。
• 低并行性： 与训练过程相比，推理过程的并行性较低，因为每个token的生成都依赖于前一个token。
• 注意力机制的二次复杂度： 注意力机制的计算量与输入序列长度的平方成正比，这会随着序列长度的增长而显著增加计算成本。

分割策略：高效利用硬件资源

为了解决上述挑战，本文提出了一个基于模型分割的框架，旨在高效利用硬件资源，并根据应用需求选择最佳分割策略。

3.1 分割符号和通信机制

本文采用了一种基于TPU v4架构的3D torus拓扑结构的分割符号和通信机制。例如，符号 BLExyz 表示将一个逻辑形状为 BLE 的张量沿最后一个维度 E 分割成 X × Y × Z 个分区，其中 x，y 和 z 分别代表TPU v4的三个物理轴，每个芯片上的张量形状为 [B, L, E/(X × Y × Z)]。

3.2 前馈层分割策略

3.2.1 一维权重固定布局

最简单的分割策略是将每个 E × F 权重矩阵沿 E 或 F 轴分割成 nchips 个分区，每个分区在相应的芯片上与激活张量进行矩阵乘法，然后使用 all-gather 和 reduce-scatter 操作进行跨芯片聚合。这种策略在芯片数量较少时效率较高，但随着芯片数量的增加，通信成本会成为瓶颈。

3.2.2 二维权重固定布局

为了提高通信效率，可以将每个 E × F 权重矩阵沿 E 和 F 轴进行二维分割，使每个分区近似为正方形。这种策略被称为二维权重固定布局。它可以有效减少通信成本，因为我们可以交替地在两个轴上进行激活张量的聚合，从而避免在每个矩阵乘法过程中都进行全量复制。

3.2.3 权重收集布局

当批处理大小和序列长度较大时，激活张量的尺寸可能会超过权重张量，此时可以将激活张量固定在每个芯片上，并将权重张量在芯片之间进行传输。这种策略被称为权重收集布局。

实验结果和结论

本文对PaLM系列巨型语言模型进行了实验验证，结果表明：

• 通过合理的模型分割策略，可以有效提高推理效率，降低延迟和成本。
• 多查询注意力机制可以有效减少内存占用，从而提高批处理大小，进而提升吞吐量。
• 在64个TPU v4芯片上，PaLM 540B模型可以实现29ms/token的低延迟生成速度，以及76%的模型FLOPS利用率，同时支持2048个token的上下文长度。

总而言之，本文提出的模型分割和优化策略为高效部署巨型语言模型提供了重要的参考，并为进一步提升推理效率提供了新的思路。

参考文献

• Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. arXiv preprint arXiv:2005.14165.
• Kaplan, J., McCandlish, S., Henighan, T., Brown, T., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., Amodei, D., et al. (2020). Scaling laws for neural language models. arXiv preprint arXiv:2001.08202.
• Rae, J. W., Borgeaud, S., Cai, T., Olah, C., Leike, J., Allen, L., Jeffery, S., Rosenthal, S., Ganguli, S., Molloy, I., et al. (2021). Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11400.
• Hoffmann, J., Habib, M., Lu, Y., Goyal, N., Zhang, X., Khandelwal, U., Das, A., Lee, K., Mishra, N., Gruslys, A., et al. (2022). Training language models with tens of trillions of parameters. arXiv preprint arXiv:2203.15556.
• Chowdhery, A., Bhatia, S., Mishra, N., Gruslys, A., Rajbhandari, S., Kumar, A., Leike, J., Allen, L., Rosenthal, S., Ganguli, S., et al. (2022). Scaling language models to 540 billion parameters. arXiv preprint arXiv:2203.15556.
• Smith, T., Zambaldi, V., Sukhbaatar, S., Raffel, C., Dhariwal, P., Leike, J., Allen, L., Rosenthal, S., Ganguli, S., Molloy, I., et al. (2022). Training language models with tens of trillions of parameters. arXiv preprint arXiv:2203.15556.
• Thoppilan, R., Sukhbaatar, S., He, J., Lee, K., Mishra, N., Gruslys, A., Rajbhandari, S., Kumar, A., Leike, J., Allen, L., et al. (2022). Scaling language models to 540 billion parameters. arXiv preprint arXiv:2203.15556.
• Sukhbaatar, S., Szlam, A., Weston, J., and Fergus, R. (2019). End-to-end efficient language modeling with data-parallel distributed attention. arXiv preprint arXiv:1907.04020.
• Choromanski, K., Rowland, M., So, A., Khan, M. E., Ballard, A., and Recht, B. (2020). Rethinking attention with performers. arXiv preprint arXiv:2009.13821.
• Dao, T., Guu, K., Lee, K., Tung, H. W., Pasupat, P., and Chang, M. W. (2022). Sparsity in deep learning: Overcoming the memory wall. arXiv preprint arXiv:2203.14722.
• Zheng, S., Li, Y., Yu, Y., Zhang, Z., and Liu, Z. (2022). Efficient large-scale language model inference on tpu v4 pods. arXiv preprint arXiv:2205.07354.
• Xu, B., Zhang, Z., Li, Y., Yu, Y., and Liu, Z. (2021). Efficient large-scale language model inference on tpus. arXiv preprint arXiv:2104.04420.
• Clarke, L., Glover, R., and MPI Forum (1994). MPI: A message-passing interface standard. Journal of Parallel and Distributed Computing, 22(1), 6-21.
• Rajbhandari, S., Rasheed, A., Madaan, A., Kumar, A., and Ganguli, S. (2020). Efficient large-scale language model training on tpus. arXiv preprint arXiv:2006.16668.
• Shoeybi, M., Patel, M., Goldfarb, C., Fevzi, B., Lee, J., Tran, L., and Parmar, N. (2019). Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053.

在现代的深度学习领域，如何优化模型性能成为了一个热门话题。许多人会依赖于一些曾经有效的小技巧，比如“使用就地操作！”、“将梯度设置为 None！”、“安装 PyTorch 1.10.0 而不是 1.10.1！”等。虽然这些方法有时能带来显著的性能提升，但从基本原理出发进行分析，往往能更系统地解决问题。

理解深度学习系统的三大组成部分

深度学习系统的性能可以分解为以下三个主要组件：

1. 计算（Compute）：在 GPU 上进行实际浮点运算（FLOPS）的时间。
2. 内存（Memory）：在 GPU 内部传输张量的时间。
3. 开销（Overhead）：其他所有时间，比如 Python 解释器运行和 CUDA 内核启动等。

理解自己深度学习系统的性能瓶颈所在，可以帮助我们有针对性地进行优化。例如，如果你的系统主要花费时间在内存传输上（即内存带宽受限），那么增加 GPU 的计算能力并不会有帮助。相反，如果你主要时间都在进行大量的矩阵乘法运算（即计算受限），那么减少开销的优化也无济于事。

接下来，我们将分别讨论计算、内存带宽和开销这三个组件。

计算：如何最大化 GPU 的 FLOPS

为了充分利用 GPU 的计算能力，我们需要尽量减少在其他部分花费的时间。GPU 的 FLOPS 越高，我们的计算效率就越高。然而，计算能力增长的速度远快于内存带宽的增长速度，使得实现高效计算变得更加困难。

例如，现代机器学习加速器（如 Nvidia 的 Tensor Cores）主要针对矩阵乘法进行了优化。如果你的计算主要不是矩阵乘法，那么你将无法完全利用 GPU 的计算能力。尽管如此，其他操作（如层归一化或激活函数）所需的 FLOPS 相对较少，因此对整体性能影响不大。

内存带宽：数据传输的成本

内存带宽成本是指在系统中移动数据的成本。这包括从 CPU 到 GPU、一个节点到另一个节点，甚至从 CUDA 全局内存到 CUDA 共享内存的传输成本。频繁的数据传输会占用大量时间，导致计算资源无法充分利用。

举个例子，当我们执行一个简单的操作如 torch.cos 时，我们将数据从存储单元移到计算单元，进行计算后再将数据移回存储单元。这种频繁的数据传输非常耗时，因此我们需要尽量减少这种操作。

操作融合：减少内存开销的关键

操作融合（Operator Fusion）是深度学习编译器中最重要的优化之一。它的基本思想是将多个计算操作合并为一个，从而减少数据在全局内存和计算单元之间的传输次数。例如，执行 x.cos().cos() 时，通常需要进行四次全局读取和写入，而通过操作融合，我们只需两次全局内存读取和写入。

这种优化在执行大规模计算时尤其重要，可以显著减少内存带宽成本，提高整体性能。

开销：代码执行之外的时间

开销指的是代码在执行实际计算和传输数据之外所花费的时间。现代 GPU 的计算速度非常快，而 Python 解释器相对较慢，因此在 Python 解释器中花费的时间会导致 GPU 资源闲置。

如何识别和减少开销

要识别是否受到开销限制，可以通过增加数据量来观察运行时间的变化。如果增加数据量后运行时间没有成比例增加，那么系统可能受到开销限制。使用 PyTorch 的分析工具也可以帮助识别 CPU 和 GPU 之间的配合情况，从而找到优化方向。

结论

要提升深度学习系统的性能，关键是识别模型的性能瓶颈，并有针对性地进行优化。以下是不同性能瓶颈的优化策略：

性能瓶颈	可能的解决方案
开销受限	使用 tracing，操作融合，避免使用 Python，采用真正的 JIT
内存带宽受限	操作融合
计算受限	使用 Tensor Cores，升级 GPU 硬件

理解这些基本原理和优化策略，可以帮助我们更有效地提升深度学习模型的性能。

---

参考文献：

• Horace He, Making Deep Learning Go Brrrr From First Principles