在 Cloud TPU 上使用 Hex-LLM 高级容器来提供开放式模型

Hex-LLM 是一种使用 XLA 供给数据的高效大语言模型 (LLM)，是专为 Cloud TPU 硬件设计和优化的 Vertex AI LLM 数据供给框架。Hex-LLM 将 LLM 数据供给技术（例如连续批处理和 PagedAttention）与专门针对 XLA 和 Cloud TPU 定制的 Vertex AI 优化结合起来。它是 Cloud TPU 上适用于开源模型的低费用、高效 LLM 数据供给服务。

可在 Model Garden 中通过模型 Playground、一键式部署和笔记本来使用 Hex-LLM。

特性

Hex-LLM 基于开源项目，并融入 Google 自己针对 XLA 和 Cloud TPU 进行的优化。Hex-LLM 在为常用的 LLM 供给数据时可实现高吞吐量和低延迟时间。

Hex-LLM 包含以下优化内容：

  • 基于词元的连续批处理算法，可帮助确保模型在处理大量并发请求时充分利用硬件。
  • 完全重写了针对 XLA 优化的注意力内核。
  • 可灵活组合的数据并行处理和张量并行处理策略，提供高度优化的权重分片方法，可在多个 Cloud TPU 芯片上高效运行 LLM。

Hex-LLM 支持各种密集和稀疏 LLM：

  • Gemma 2B 和 7B
  • Gemma 2 9B 和 27B
  • Llama 2 7B、13B 和 70B
  • Llama 3 8B 和 70B
  • Mistral 7B 和 Mixtral 8x7B

Hex-LLM 还有诸多特性，例如：

  • Hex-LLM 被封装在单个容器中。Hex-LLM 将 API 服务器、推理引擎和支持的模型打包到单个 Docker 映像中进行部署。
  • 与 Hugging Face 模型格式兼容。Hex-LLM 可以从本地磁盘、Hugging Face Hub 和 Cloud Storage 存储桶加载 Hugging Face 模型。
  • 使用 bitsandbytes 和 AWQ 进行量化处理。
  • 动态 LoRA 加载。Hex-LLM 能够在供给数据期间通过读取请求参数来加载 LoRA 权重。

开始使用 Model Garden

Hex-LLM Cloud TPU 数据供给容器已集成到 Model Garden 中。您可以通过该数据供给技术提供的 Playground、一键式部署和 Colab Enterprise 笔记本示例在各种模型中使用该技术。

使用 Playground

Model Garden Playground 是一个预部署的 Vertex AI 端点，可通过在模型卡片中发送请求访问该端点。

  1. 输入提示，并视情况为请求添加参数。

  2. 点击提交以快速获取模型回答。

通过 Gemma 试用一下吧！

使用一键式部署

您可以使用模型卡片部署提供 Hex-LLM 的自定义 Vertex AI 端点。

  1. 前往模型卡片页面，然后点击部署。

  2. 为要使用的模型变体选择 Cloud TPU v5e 机器类型进行部署。

  3. 点击底部的部署以开始部署流程。您会收到两条电子邮件通知：在模型上传完成时以及在端点准备就绪时。

使用 Colab Enterprise 笔记本

为了实现灵活的自定义内容，您可以使用 Colab Enterprise 笔记本示例通过 Vertex AI SDK for Python 部署提供 Hex-LLM 的 Vertex AI 端点。

  1. 前往模型卡片页面，然后点击打开笔记本。

  2. 选择 Vertex Serving 笔记本。该笔记本随即会在 Colab Enterprise 中打开。

  3. 运行该笔记本，以使用 Hex-LLM 部署模型，并将预测请求发送到端点。部署的代码段如下所示：

hexllm_args = [
    f"--model=google/gemma-2-9b-it",
    f"--tensor_parallel_size=4",
    f"--hbm_utilization_factor=0.8",
    f"--max_running_seqs=512",
]
hexllm_envs = {
    "PJRT_DEVICE": "TPU",
    "MODEL_ID": "google/gemma-2-9b-it",
    "DEPLOY_SOURCE": "notebook",
}
model = aiplatform.Model.upload(
    display_name="gemma-2-9b-it",
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=[
        "python", "-m", "hex_llm.server.api_server"
    ],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=hexllm_envs,
    serving_container_shared_memory_size_mb=(16 * 1024),
    serving_container_deployment_timeout=7200,
)

endpoint = aiplatform.Endpoint.create(display_name="gemma-2-9b-it-endpoint")
model.deploy(
    endpoint=endpoint,
    machine_type="ct5lp-hightpu-4t",
    deploy_request_timeout=1800,
    service_account="<your-service-account>",
    min_replica_count=1,
    max_replica_count=1,
)

您可以修改以下 Hex-LLM 服务器启动参数来自定义数据供给：

  • --model：要加载的模型。您可以指定 Hugging Face 模型 ID、本地绝对路径或 Cloud Storage 存储桶路径。
  • --tokenizer：要加载的词元化器。您可以指定 Hugging Face 模型 ID、本地绝对路径或 Cloud Storage 存储桶路径。默认值与 --model 的值相同。
  • --enable_jit：指定是否启用 JIT 模式。默认值为 True。
  • --enable_lora：指定是否启用 LoRA 加载模式。默认值为 False。
  • --max_lora_rank：LoRA 适配器支持的 LoRA 秩上限（在请求中定义）。默认值为 16。
  • --data_parallel_size：数据并行副本的数量。默认值为 1。
  • --tensor_parallel_size：张量并行副本的数量。默认值为 1。
  • --max_running_seqs：服务器可以并行处理的请求数上限。此参数值越大，服务器可以实现的吞吐量越高，但可能对延迟时间有一定影响。默认值为 256。
  • --hbm_utilization_factor：加载模型权重后，可为 KV 缓存分配的免费 Cloud TPU HBM 所占的百分比。将此参数设置为较小的值可以有效地防止 Cloud TPU HBM 内存不足。默认值为 0.9。
  • --seed：用于初始化所有随机数生成器的种子。更改此参数可能会影响针对同一提示生成的输出。默认值为 0。

申请 Cloud TPU 配额

在 Model Garden 中，您的默认配额为 us-west1 区域中的四个 Cloud TPU v5e 芯片。此配额适用于一键式部署和 Colab Enterprise 笔记本部署。如需申请更多配额，请参阅申请更高配额。