自动扩缩使用 TPU 的大语言模型 (LLM) 推理工作负载的最佳实践


本最佳实践指南向您展示了可用的指标,以及如何选择合适的指标来为 GKE 上使用 TPU 的单主机 JetStream 推理工作负载设置 Pod 横向自动扩缩器 (HPA)。HPA 是确保模型服务器随负载进行适当扩缩的一种高效方法。微调 HPA 设置是使预配的硬件费用与流量需求保持一致以实现推理服务器性能目标的主要方式。

如需查看实现这些最佳实践的示例,请参阅通过 GKE 为 TPU 上的 LLM 工作负载配置自动扩缩

目标

本指南适用于对通过 Kubernetes 来优化使用 TPU 的单主机 JetStream 工作负载感兴趣的生成式 AI 客户、新的或现有 GKE 用户、机器学习工程师和 LLMOps (DevOps) 工程师。

阅读本指南后,您应该能够:

  • 了解适用于单主机 JetStream 推理的关键自动扩缩指标。
  • 在考虑要自动扩缩的指标时,请了解整体利弊。

适用于 JetStream 推理的自动扩缩指标概览

我们建议您使用以下指标:

服务器指标

与许多其他 LLM 推理服务器一样,JetStream 提供了性能指标。GKE 可以根据这些服务器级指标简化 JetStream 的监控和自动扩缩。如需配置任何自动扩缩,您必须先了解 JetStreams 请求流水线如何影响关键性能指标。所有传入请求都会经过预填充队列、传输队列和生成队列,然后成为解码请求。JetStream 会陆续接受解码请求,并使用固定数量的解码线程并发处理这些请求。在给定时间点,处理解码请求所占用的解码引擎百分比以 jetstream_slots_used_percentage 指标进行测量。

扩缩单主机 JetStream 时,这会对延迟时间和吞吐量产生两个影响:

  • 如果传入请求的速率低于解码槽可以处理请求的速率,则请求不会积压在队列中。如果 JetStream 没有积压,则吞吐量会与传入请求的速率成正比。延迟时间基本上保持不变,但会随着并发解码请求的数量而略微成比例增加,因为新接受的解码请求会中断解码。
  • 一旦传入请求的速率超过解码槽可以处理请求的速率,请求便会积压在队列中。如果 JetStream 存在积压,则请求延迟时间会显著增加,并且与积压请求数量成正比,而吞吐量会保持不变。

以下服务器指标与这些相关性能指标的相关性最强,我们建议您使用这些指标进行自动扩缩:

这些指标通常能够灵活应对性能和流量波动,因此是在各种 TPU 硬件设置中自动扩缩的可靠起点。

TPU 指标

鉴于 LLM 服务会在内存而非计算方面形成瓶颈,我们建议您根据内存用量而不是其他 TPU 指标来扩缩 JetStream,因为它能够最好地反映硬件的资源利用率。如需了解详情,请参阅相关的最佳实践部分。

选择自动扩缩指标的注意事项

请考虑以下注意事项和最佳实践,以选择在 GKE 上自动扩缩的最佳指标,以满足推理工作负载性能目标。

最佳实践:使用预填充积压(队列)大小在一定的目标延迟时间阈值范围内最大限度地提高吞吐量并最大限度地降低费用

在优化吞吐量和费用时,以及使用模型服务器的每设备批次大小的最大吞吐量可以实现延迟时间目标时,我们建议进行预填充队列大小自动扩缩。

预填充队列大小与请求延迟时间直接相关。传入请求在处理之前会在模型服务器预填充队列中排队,因此此队列时间会增加整体延迟时间。队列大小是负载峰值的敏感指标,因为增加的负载会快速填充队列。

基于预填充队列大小的自动扩缩可在负载下扩容,并在队列为空时缩容,从而最大限度地缩短队列时间。这种方法相对容易实现,因为队列大小与请求大小、模型或硬件无关。

如果您希望最大限度地提高每个模型服务器副本的吞吐量,请考虑重点关注预填充队列大小。预填充队列大小会跟踪待处理(而不是正在处理处理)的请求。JetStream 使用连续批处理,这样可以最大限度地提高并发请求数量,并在批次空间可用时使队列保持较低水平。当批次空间有限时,队列会明显增大,因此请使用增长点作为启动扩容的信号。通过将队列大小自动扩缩与优化的批次吞吐量相结合,您可以最大限度地提高请求吞吐量。

为 HPA 确定最佳队列大小阈值

如需选择正确的队列大小阈值,请从 3-5 之间的值开始,然后逐渐增加该值,直到请求达到偏好的延迟时间。可使用 locust-load-inference 工具进行测试。对于低于 10 的阈值,请微调 HPA 扩容设置以处理流量峰值。

您还可以创建 Cloud Monitoring 自定义信息中心来直观呈现指标行为。

限制

请注意 HPA 容忍,这默认为目标值周围的 0.1 无操作范围,用于抑制振荡。

预填充队列大小不会直接控制并发请求数量,因此其阈值无法保证延迟时间低于每设备批次大小允许的延迟时间。如需解决此问题,您可以手动减小每设备批次大小,或根据批次大小自动扩缩

最佳实践:使用 slots_used 百分比可达到比队列大小更低的目标延迟时间阈值

如果您的工作负载对延迟时间敏感,并且基于队列的扩缩不够快,无法满足您的要求,我们建议您选择基于 slots_used 的自动扩缩。

根据 slot_used 自动扩缩可确保工作负载吞吐量能够扩容以最大限度地提高同时并行处理的请求数量,并在并行处理的请求数量减少时缩容。这会对延迟时间产生两个影响。首先,由于基于 slots_used 的自动扩缩是为了确保每个传入请求都有一个槽,因此阈值设置为 1 的接近程度会与请求花费时间排队的可能性相对应,从而具有更长的延迟时间。其次,更大的批次大小会提高吞吐量,但也会增加延迟时间,因为在连续批处理模型服务器中,某些请求的预填充阶段会中断其他请求的解码阶段。您可以监控批次大小模式,并使用自动扩缩在负载峰值期间最大限度地减少并发请求数量。

如果队列大小已满足您的延迟时间目标,请优先考虑自动扩缩。 这样可以最大限度地提高吞吐量和成本效益。不过,slots_used 对于于延迟时间敏感型工作负载很有用。

我们还建议您先将每设备批次大小调优为适当的值,然后再探索基于 slots_used 的自动扩缩。(可选)您也可以将此设置与基于队列的自动扩缩结合使用。

为 HPA 确定最佳 slots_used 百分比阈值

如需选择合适的批次大小阈值,请通过实验增加服务器上的负载,并观察批次大小达到峰值的情况。我们还建议使用 locust-load-inference 工具进行测试。确定会最大限度提高内存用量的每设备批次大小后,您就可以配置 slots_used 百分比目标值。将初始目标值设置为略低于 1,然后降低该值,直到达到偏好的延迟时间。

您还可以创建 Cloud Monitoring 自定义信息中心来直观呈现指标行为。

限制

请注意 HPA 容忍,这默认为目标值周围的 0.1 无操作范围,用于抑制振荡。

根据 slot_used 百分比自动扩缩虽然有助于控制延迟时间,但存在限制。不同的请求大小和硬件限制使得找到合适的 slot_used 百分比阈值变得困难。如果某个扩缩规则尝试将平均 slots_used 百分比保持在 100% 以下,则表示该扩缩规则尝试保持可用槽数不为零。这些可用槽对应于未使用的可用吞吐量,如果您想要充分利用可用 TPU,这并不是理想情况。

最佳实践:使用 TPU 高带宽内存 (HBM) 用量最大限度地提高硬件利用率

TPU 高带宽内存 (HBM) 用量与硬件利用率具有最直接的对应关系,即使与系统指标相比也是如此,因为它们不考虑请求大小。虽然根据队列大小进行扩缩可以更好地最大限度提高硬件利用率,但代价是延迟时间会增加。如果您更倾向于依赖系统指标而非服务器指标,我们建议您使用 HBM 用量作为根据 slots_used 自动扩缩的最佳替代方案,因为二者都与吞吐量相对应。如需详细了解 TPU 内存,请参阅 TPU 的工作方式

将批次大小增加到超过最佳点可提高吞吐量,但也会增加出现内存不足 (OOM) 错误的风险。我们强烈建议您根据 HBM 指标进行扩缩,以平衡吞吐量和稳定性。我们建议不要同时根据预填充队列大小和 HBM 用量进行扩缩,因为随着负载增加,HBM 用量将会增加,并首先触发扩缩。

为 HPA 确定最佳 TPU HBM 用量阈值

在选择内存用量阈值之前,我们建议将每设备批次大小设置为在预期最大负载下运行时可最大限度地使用内存的值。请注意,该值需要通过实验方式确定,在很大程度上取决于总 HBM 以及预期的提示和回答长度。我们建议您使用 locust-load-inference 工具进行实验和测试。

与每设备批次大小类似,设置阈值以最大限度地提高 TPU 内存利用率,同时最大限度地降低 OOM 行为的风险。

您还可以创建 Cloud Monitoring 自定义信息中心来直观呈现指标行为。

限制

有两个注意事项会削弱延迟时间和吞吐量与 HBM 用量对应的强度,即 HBM 用量的波动性以及通常 TPU 指标的较低采样率。此外,虽然 HBM 用量与延迟时间之间仍然存在对应关系,但 HBM 用量的增加对延迟时间的影响远远小于排队请求数量的增加。

最佳实践:优化 HPA 配置

我们建议您设置以下 HPA 配置选项:

  • 稳定时长:使用此 HPA 配置选项可防止因指标波动而导致副本数量快速变化。默认值为 5 分钟表示缩容(避免过早缩容),0 表示扩容(确保响应速度)。请根据工作负载波动性和您偏好的响应速度调整该值。
  • 扩缩政策:使用此 HPA 配置选项可微调扩容和缩容行为。您可以设置“Pod”政策限制来指定每个时间单位内发生变化的副本的绝对数量,并设置“百分比”政策限制来按百分比变化进行指定。

如需详细了解这些选项,请参阅开源 Kubernetes 文档中的 Pod 横向自动扩缩

后续步骤