使用 Google Kubernetes Engine (GKE) 上的 GPU 优化大语言模型推理的最佳实践

Google Kubernetes Engine (GKE) 可对大语言模型 (LLM) 推理进行精细控制,同时实现最佳性能和成本效益。本指南介绍了通过 vLLM文本生成推理 (TGI) 服务框架,使用 GKE 上的 GPU 优化开放 LLM 的推理和服务的最佳实践。

如需查看所有最佳实践的汇总核对清单,请参阅核对清单摘要

如需全面了解所有 GKE 最佳实践,请参阅 GKE 最佳实践

目标

本指南适用于关注使用 GPU 和 Kubernetes 来优化 LLM 工作负载的生成式 AI 客户、新 GKE 用户或现有 GKE 用户、机器学习工程师和 LLMOps (DevOps) 工程师。

学完此指南后,您将能够:

  • 选择训练后 LLM 优化技术,包括量化、张量并行处理和内存优化。
  • 在考虑这些优化技术时,需要权衡利弊。
  • 使用启用了优化设置的 vLLM 或 TGI 等服务框架将开放 LLM 模型部署到 GKE。

LLM 服务优化技术概览

与非 AI 工作负载不同,LLM 工作负载依赖于矩阵乘法运算,因此通常表现出更高的延迟时间和更低的吞吐量。如需增强 LLM 推理性能,您可以使用专用硬件加速器(例如 GPU 和 TPU)和经过优化的服务框架。

您可以应用以下一种或多种最佳实践来缩短 LLM 工作负载延迟时间,同时提高吞吐量和成本效益:

本指南中的示例将 Gemma 7B LLM 与 vLLM 或 TGI 服务框架结合使用,以应用这些最佳实践;不过,所述的概念和功能适用于大多数流行的开放 LLM。

准备工作

在试用本指南中的示例之前,请完成以下前提任务:

  1. 请按照以下指南中的说明访问 Gemma 模型、准备环境以及创建和配置 Google Cloud 资源:

    请务必将 Hugging Face 访问令牌保存到您的 Kubernetes Secret 中。

  2. https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/ 示例代码库克隆到您的本地开发环境。

  3. 将您的工作目录更改为 /kubernetes-engine-samples/ai-ml/llm-serving-gemma/

最佳实践:量化

量化是一种与有损图像压缩类似的技术,通过以低精度格式(8 位或 4 位)表示权重来缩减模型大小,从而降低内存需求。但是,与图像压缩一样,量化也存在权衡:模型大小减小会导致准确率降低。

存在各种量化方法,每种方法都有其独特的优点和缺点。有些量化方法(例如 AWQ 和 GPTQ)需要预量化,并且在 Hugging FaceKaggle 等平台上提供。例如,如果您对 Llama-2 13B 模型应用 GPTQ,并对 Gemma 7B 模型应用 AWQ,则可以在单个 L4 GPU 上(而不是两个 L4 GPU 上)提供模型,而无需量化。

您还可以使用 AutoAWQAutoGPTQ 等工具执行量化。这些方法可以缩短延迟时间并提高吞吐量。相比之下,使用 EETQbitsandbytes 库进行量化的技术不需要预量化模型,因此在没有预量化版本时可以成为合适的选择。

要使用的最佳量化技术取决于您的特定目标,以及该技术与您要使用的服务框架的兼容性。如需了解详情,请参阅 Hugging Face 中的量化指南

选择其中一个标签页可查看使用 TGI 或 vLLM 框架应用量化的示例:

TGI

GKE 通过 TGI 支持以下量化选项:

  • awq
  • gptq
  • eetq
  • bitsandbytes
  • bitsandbytes-nf4
  • bitsandbytes-fp4

AWQ 和 GPTQ 量化方法需要使用预量化模型,而 EETQ 和 bitsandbytes 量化可应用于任何模型。如需详细了解这些选项,请参阅这篇 Hugging Face 文章

如需使用量化,请在启动模型服务器时设置 -–quantize 参数。

以下代码段展示了如何在 GKE 上使用 TGI 通过 bitsandbytes 量化进行 Gemma 7B 优化。

args:
- --model-id=$(MODEL_ID)
- --num-shard=2
- --quantize=bitsandbytes

如需应用此配置,请使用以下命令:

kubectl apply -f tgi/tgi-7b-bitsandbytes.yaml

vLLM

GKE 通过 vLLM 支持以下量化选项:

如需将模型量化与 vLLM 搭配使用,必须预先对模型进行量化。若启动运行时,请设置 –quantization 参数。

以下代码段说明了如何使用 GKE 上的 vLLM 通过 awq 量化进行 Gemma 7B 模型优化:

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --quantization=awq
env:
- name: MODEL_ID
  value: google/gemma-7b-AWQ
resources:
  requests:
    nvidia.com/gpu: 1
  limits:
    nvidia.com/gpu: 1

如需应用此配置,请使用以下命令:

kubectl apply -f vllm/vllm-7b-awq.yaml

通过使用 KV 缓存量化来缩短延迟时间

您可以使用 FP8 E5M2 KV 缓存量化来显著减少 KV 缓存内存占用并缩短延迟时间,尤其是对于较大批次大小。但是,这会降低推断准确率。

如需启用 FP8 E5M2 KV 缓存量化,请设置参数 --kv-cache-dtype fp8_e5m2

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --kv-cache-dtype=fp8_e5m2
- --max-model-len=1200
resources:
  requests:
    cpu: "2"
    memory: "25Gi"
    ephemeral-storage: "25Gi"
    nvidia.com/gpu: 1
  limits:
    cpu: "2"
    memory: "25Gi"
    ephemeral-storage: "25Gi"
    nvidia.com/gpu: 1

如需应用此配置,请使用以下命令:

kubectl apply -f vllm/vllm-7b-kvcache.yaml

最佳做法:张量并行处理

张量并行处理是一种将计算负载分布到多个 GPU 上的技术,在运行超出单个 GPU 内存容量的大型模型时,这一点至关重要。这种方法的性价比更高,因为您可以使用多个价格实惠的 GPU,而不是单个价格昂贵的 GPU。它还可以提高模型推理吞吐量。张量并行处理利用了张量运算可以在较小的数据块上独立执行的事实。

如需详细了解此技术,请参阅 Hugging Face 的张量并行处理指南

选择其中一个标签页以查看使用 TGI 或 vLLM 框架应用张量并行处理的示例:

TGI

使用 TGI 时,默认情况下,服务运行时将使用 Pod 可用的所有 GPU。如需设置要使用的 GPU 数量,您可以指定 --num-shard 参数并将 GPU 数量的值设为该值。

如需查看支持张量并行处理的模型列表,请参阅 Hugging Face 文档。

以下代码段展示了如何使用张量并行性和两个 L4 GPU 优化 Gemma 7B 指令调优模型:

args:
- --model-id=$(MODEL_ID)
- --num-shard=2

如需应用此配置,请使用以下命令:

kubectl apply -f tgi/tgi-7b-it-tensorparallelism.yaml

在 GKE Autopilot 集群中,运行此命令会创建一个 Pod,其最低资源要求为 21 个 vCPU 和 78 GiB 内存。

vLLM

vLLM 支持分布式张量并行推理。如果有多个 GPU,vLLM 会默认启用该功能。

以下代码段介绍了如何使用张量并行处理和两个 L4 GPU 优化 Gemma 7B 指令调优模型:

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=2

如需应用此配置,请使用以下命令:

kubectl apply -f vllm/vllm-7b-it-tensorparallelism.yaml

在 GKE Autopilot 集群中,运行此命令会创建一个 Pod,其最低资源要求为 21 个 vCPU 和 78 GiB 内存。

最佳实践:模型内存优化

优化 LLM 的内存用量对于高效推理至关重要。本部分介绍了注意力层优化策略,例如分页注意力和闪存注意力。这些策略可提高内存效率,支持更长的输入序列并缩短 GPU 空闲时间。本部分还介绍了如何调整模型输入和输出大小以适应内存限制,以及如何针对特定服务框架进行优化。

注意力层优化

自注意力层可让模型在语言处理任务中理解上下文,因为字词的含义可能会因上下文而异。但是,这些层会在 GPU vRAM 中存储输入令牌权重、键 (K) 和值 (V)。因此,随着输入序列的延长,这会导致大小和计算时间的二次增长。

在处理较长的输入序列时,使用 KV 缓存尤其有用,其中自注意力的开销可能非常大。这种优化方法可将计算处理降至线性复杂性。

用于优化 LLM 中注意力机制的具体技术包括:

  • 分页注意力分页注意力通过使用与操作系统虚拟内存类似的分页技术,改进了大模型和长输入序列的内存管理。这可有效地减少 KV 缓存中的碎片和重复,支持更长的输入序列,而不会耗尽 GPU 内存。
  • 闪存注意力闪存注意力通过最大限度地减少令牌生成期间 GPU RAM 和 L1 缓存之间的数据传输,进而减少 GPU 内存瓶颈。这样可以消除计算核心的空闲时间,显著提高 GPU 的推理和训练性能。

模型输入和输出大小调整

内存要求取决于输入和输出大小。输出越长、上下文越多,需要的资源就越多;而输出越短、上下文越少,则可以通过使用体积更小、价格更低的 GPU 来节省费用。

选择其中一个标签页以查看在 TGI 或 vLLM 框架中调整模型输入和输出内存要求的示例:

TGI

TGI 服务运行时会在启动期间检查内存要求,如果可能的最大模型内存占用量不适合可用的 GPU 内存,则不会启动。此检查可消除内存密集型工作负载上发生的内存不足 (OOM) 崩溃。

GKE 支持以下 TGI 参数,以优化模型内存要求:

以下代码段说明了如何使用单个 L4 GPU 和参数设置 --max-total-tokens=3072, --max-batch-prefill-tokens=512, --max-input-length=512 来提供 Gemma 7B 指令调优型模型:

args:
- --model-id=$(MODEL_ID)
- --num-shard=1
- --max-total-tokens=3072 
- --max-batch-prefill-tokens=512
- --max-input-length=512
env:
- name: MODEL_ID
  value: google/gemma-7b

如需应用此配置,请使用以下命令:

kubectl apply -f tgi/tgi-7b-token.yaml

vLLM

在 vLLM 中,配置模型的上下文长度,这直接影响 KV 缓存大小和 GPU RAM 要求。上下文长度越短,可以使用价格更实惠的 GPU。默认值为模型接受的词法单元数量上限。如果需要,请使用 --max-model-len MAX_MODEL_LEN 限制上下文长度上限。

例如,Gemma 7B 指令调优模型的默认上下文长度为 8192,超过了单个 NVIDIA L4 GPU 的内存容量。如需在 L4 上部署,请将 --max-model-len 设置为小于 640 的值,以限制提示和输出的总长度。此调整支持在单个 L4 GPU 上运行模型,尽管该模型的默认上下文长度较长。

如需使用修改后的令牌限制进行部署,请使用以下代码段:

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --max-model-len=600

如需应用此配置,请使用以下命令:

kubectl apply -f vllm/vllm-7b-token.yaml

核对清单摘要

优化目标 练习
延迟时间
  • 优先使用功能强大的 GPU:请考虑升级到具有更高计算能力和内存 I/O 吞吐量的 GPU。
  • 探索量化:AWQ 等技术可以缩短延迟时间,但要注意潜在的准确率权衡。
吞吐量
  • 横向扩缩:增加提供服务的副本 (Pod) 的数量,以分散工作负载。
  • 使用张量并行处理:对于超出单个 GPU 容量的大型模型,请使用张量并行功能将计算分布到多个 GPU 上。对于较小的模型,请考虑使用张量并行处理值为“1”的多个副本,以避免开销。
  • 批量请求和量化:组合请求并探索可保持可接受准确度的量化技术。
成本效益
  • 选择较小的模型:选择系列中适合您的资源限制和预算的模型。
  • 应用量化:使用量化来减少内存需求,尤其是在处理较大的模型时。
  • 限制上下文长度:限制上下文长度,以进一步减少内存使用量,并在更小、更经济实惠的 GPU 上执行。

后续步骤