动机

服务大型模型通常会导致内存瓶颈，例如可怕的 CUDA 内存不足 错误。为了解决这个问题，主要有两种解决方案

  1. 降低精度 – 利用 FP8 和更低位量化方法可以减少内存使用。然而，这种方法可能会影响准确性和可扩展性，并且随着模型增长到数千亿参数以上，这种方法本身并不足够。
  2. 分布式推理 – 将模型计算分布在多个 GPU 或节点上可以提高可扩展性和效率。 这就是张量并行和流水线并行等分布式架构发挥作用的地方。

vLLM 架构和大型语言模型推理挑战

与训练相比，LLM 推理带来了独特的挑战

  • 与专注于具有已知静态形状的吞吐量的训练不同，推理需要低延迟和动态工作负载处理。
  • 推理工作负载必须有效管理 KV 缓存、推测解码和预填充到解码的转换。
  • 大型模型通常超出单 GPU 容量，需要先进的 并行化策略。

为了解决这些问题，vLLM 提供了

  • 张量并行，用于在节点内的多个 GPU 之间分片每个模型层。
  • 流水线并行，用于跨多个节点分发模型层的连续部分。
  • 优化的通信内核和控制平面架构，以最大限度地减少 CPU 开销并最大化 GPU 利用率。

vLLM 中的 GPU 并行技术

张量并行

问题：模型超出单 GPU 容量

随着模型增长，单个 GPU 无法容纳它们，因此需要多 GPU 策略。张量并行在 GPU 之间分片模型权重，从而允许并发计算以实现更低的延迟和增强的可扩展性。

这种方法最初是为 Megatron-LM (Shoeybi et al., 2019) 中的训练而开发的，已在 vLLM 中针对推理工作负载进行了调整和优化。

张量并行依赖于两种主要技术

  1. 列并行：沿列分割权重矩阵，并在计算后连接结果。
  2. 行并行：沿行分割矩阵，并在计算后对部分结果求和。

作为一个具体的例子，让我们分解一下这种并行性如何在 Llama 模型中的 MLP（多层感知器）层中工作

  • 列并行适用于向上投影操作。
  • 逐元素激活函数（例如，SILU）对分片输出进行操作。
  • 行并行用于向下投影，并使用 all-reduce 操作来聚合最终结果。

张量并行确保推理计算分布在多个 GPU 上，从而最大限度地利用可用的内存带宽和计算。 使用时，我们可以通过有效地倍增内存带宽来获得延迟改进。 发生这种情况是因为分片模型权重允许多个 GPU 并行访问内存，从而减少单个 GPU 可能遇到的瓶颈。

来源：Sebastian Raschka, 2023 年。

然而，它需要在每个 GPU 之间使用 高带宽互连，例如 NVLink 或 InfiniBand，以最大限度地减少因增加的通信成本而产生的开销。

流水线并行

问题：模型超出多 GPU 容量

对于极其大型的模型（例如，DeepSeek R1、Llama 3.1 405B），单个节点可能不足以满足需求。 流水线并行跨节点分片模型，每个节点处理特定的连续模型层。

工作原理

  • 每个 GPU 加载和处理一组不同的层。
  • 发送/接收操作： 中间激活在 GPU 之间传输，随着计算的进行。

与张量并行相比，这导致更低的通信开销， 因为数据传输在每个流水线阶段发生一次。

流水线并行减少了 GPU 之间的内存限制，但不会像张量并行那样从根本上减少推理延迟。 为了缓解吞吐量效率低下，vLLM 采用了先进的流水线调度，通过优化微批量执行来确保所有 GPU 保持活动状态。

结合张量并行和流水线并行

作为一般经验法则，可以这样考虑并行性的应用

  • 当互连速度较慢时，跨节点使用流水线并行，在节点内使用张量并行。
  • 如果互连高效（例如，NVLink、InfiniBand），则 张量并行可以跨节点扩展。
  • 智能地结合这两种技术可以 减少不必要的通信开销 并最大限度地提高 GPU 利用率。

性能扩展和内存效应

虽然并行化的基本原理表明线性扩展，但实际上，由于内存效应，性能改进可能是超线性的。 无论是张量并行还是流水线并行，由于 KV 缓存的可用内存超线性增加，吞吐量改进可能会以不明显的方式出现。

这种超线性扩展效应的发生是因为更大的缓存允许更大的批量大小，以便并行处理更多请求和更好的内存局部性，从而提高 GPU 利用率，超出仅添加更多计算资源所能预期的水平。 在上图中，您可以看到在 TP=1 和 TP=2 之间，我们能够将 KV 缓存块的数量增加 13.9 倍，这使我们能够观察到 3.9 倍的令牌吞吐量 - 远远超过我们预期从使用 2 个 GPU 而不是 1 个 GPU 获得的线性 2 倍。

进一步阅读

对于有兴趣深入了解影响 vLLM 设计的技术和系统的读者

结论

高效地服务大型模型需要结合 张量并行、流水线并行 和 性能优化（如 分块预填充）。 vLLM 通过利用这些技术实现可扩展的推理，同时确保跨不同硬件加速器的适应性。 随着我们继续增强 vLLM，及时了解新的发展，例如 混合专家 (MoE) 的专家并行 和 扩展的量化支持，对于优化 AI 工作负载至关重要。

参加双周办公时间，了解更多关于 LLM 推理优化和 vLLM 的信息！

致谢

感谢 Sangbin Cho (xAI) 创作了一些图表。