在 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
• Llama-3.1 8B 和 70B
• Llama-3.2 1B 和 3B
• Llama-3.3 70B
• Llama-Guard-3 1B 和 8B
• Llama-4 Scout-17B-16E
• Mistral 7B
• Mixtral 8x7B 和 8x22B
• Phi-3 mini 和 medium
• Phi-4、Phi-4 推理和推理 Plus
• Qwen-2 0.5B、1.5B 和 7B
• Qwen-2.5 0.5B、1.5B、7B、14B 和 32B

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 权重。

高级功能

Hex-LLM 支持以下高级功能：

• 多主机服务
• 分离式服务 [实验性]
• 前缀缓存
• 支持 4 位量化

多主机服务

Hex-LLM 现在支持使用多主机 TPU 切片来部署模型。借助此功能，您可以部署无法加载到单个主机 TPU 虚拟机（最多包含 8 个 v5e 核心）中的大型模型。

如需启用此功能，请在 Hex-LLM 容器参数中设置 --num_hosts，并在 Vertex AI SDK 模型上传请求中设置 --tpu_topology。以下示例展示了如何部署 Hex-LLM 容器，该容器采用 TPU 4x4 v5e 拓扑，用于提供 Llama 3.1 70B bfloat16 模型：

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Meta-Llama-3.1-70B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=16",
    "--num_hosts=4",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    tpu_topology="4x4",
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

如需查看有关如何使用多主机 TPU 拓扑部署 Hex-LLM 容器的端到端教程，请参阅 Vertex AI Model Garden - Llama 3.1（部署）笔记本。

一般来说，只需进行以下更改即可启用多主机服务：

1. 将参数 --tensor_parallel_size 设置为 TPU 拓扑中的总核心数。
2. 将参数 --num_hosts 设置为 TPU 拓扑中的主机数量。
3. 使用 Vertex AI SDK 模型上传 API 设置 --tpu_topology。

分离式服务 [实验性]

Hex-LLM 现在支持分离式服务（实验性功能）。它只能在单主机设置中启用，并且性能正在调优中。

分离式服务是一种有效的方法，可平衡每个请求的第一个 token 的时间 (TTFT) 和每个输出 token 的时间 (TPOT)，以及整体服务吞吐量。它将预填充阶段和解码阶段分离为不同的工作负载，以便它们互不干扰。此方法对于设置了严格延迟要求的场景尤其有用。

如需启用此功能，请在 Hex-LLM 容器参数中设置 --disagg_topo。 以下示例展示了如何在 TPU v5e-8 上部署 Hex-LLM 容器，以提供 Llama 3.1 8B bfloat16 模型：

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=2",
    "--disagg_topo=3,1",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

--disagg_topo 参数接受 "number_of_prefill_workers,number_of_decode_workers" 格式的字符串。在前面的示例中，它设置为 "3,1"，用于配置 3 个预填充工作器和 1 个解码工作器。每个工作器使用两个 TPU v5e 核心。

前缀缓存

前缀缓存可缩短在提示开头具有相同内容的提示的第一个token时间 (TTFT)，例如公司范围内的序言、通用系统指令和多轮对话历史记录。Hex-LLM 不会重复处理相同的输入 token，而是会保留已处理的输入 token 计算的临时缓存，以缩短 TTFT。

如需启用此功能，请在 Hex-LLM 容器参数中设置 --enable_prefix_cache_hbm。以下示例展示了如何在 TPU v5e-8 上部署 Hex-LLM 容器，以部署 Llama 3.1 8B bfloat16 模型：

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_prefix_cache_hbm",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

Hex-LLM 采用前缀缓存来优化超过一定长度（默认情况下为 512 个 token，可使用 prefill_len_padding 进行配置）的提示的性能。缓存命中以该值的增量发生，确保缓存的 token 数量始终是 prefill_len_padding 的倍数。对话完成 API 响应中 usage.prompt_tokens_details 的 cached_tokens 字段表示有多少提示token是缓存命中。

"usage": {
  "prompt_tokens": 643,
  "total_tokens": 743,
  "completion_tokens": 100,
  "prompt_tokens_details": {
    "cached_tokens": 512
  }
}

分块预填充

分块预填充将请求预填充拆分为更小的块，并将预填充和解码混合到一个批处理步骤中。Hex-LLM 实现分块预填充，以平衡第一个 token 的时间 (TTFT) 和每个输出 token 的时间 (TPOT)，并提高吞吐量。

如需启用此功能，请在 Hex-LLM 容器参数中设置 --enable_chunked_prefill。以下示例展示了如何在 TPU v5e-8 上部署 Hex-LLM 容器，以部署 Llama 3.1 8B 模型：

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_chunked_prefill",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

支持 4 位量化

量化是一种通过使用低精度数据类型（如 INT8 或 INT4）而非通常的 BF16 或 FP32 来表示权重或激活，从而降低运行推理的计算和内存成本的技术。

Hex-LLM 支持仅限权重的 INT8 量化。扩展支持包括使用 AWQ 零点量化进行量化的 INT4 权重模型。Hex-LLM 支持 Mistral、Mixtral 和 Llama 模型系列的 INT4 变体。

无需额外的标志即可提供量化模型。

开始使用 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,
)

Colab Enterprise 笔记本示例包括：

• Gemma 2 部署
• CodeGemma 部署
• Llama 3.2 部署
• Llama 3.1 部署
• Phi-3 部署
• Qwen2 部署

配置服务器参数和环境变量

您可以设置以下参数来启动 Hex-LLM 服务器。您可以根据预期使用场景和要求量身定制参数。请注意，这些参数是为一键式部署预定义的，旨在提供最简单的部署体验。如需自定义参数，您可以参考笔记本示例，并相应地设置参数。

模型

• --model：要加载的模型。您可以指定 Hugging Face 模型 ID、Cloud Storage 存储桶路径 (gs://my-bucket/my-model) 或本地路径。模型制品应遵循 Hugging Face 格式，并使用 safetensors 文件存储模型权重。Llama、Gemma 2 和 Mistral/Mixtral 支持 BitsAndBytes int8 和 AWQ 量化模型制品。
• --tokenizer：要加载的词元化器。可以是 Hugging Face 模型 ID、Cloud Storage 存储桶路径 (gs://my-bucket/my-model) 或本地路径。如果未设置此参数，则默认为 --model 的值。
• --tokenizer_mode：词元化器模式。可能的选项为 ["auto", "slow"]。默认值为 "auto"。如果将其设置为 "auto"，则会使用快速词元化器（如果可用）。慢速词元化器是用 Python 编写的，并在 Transformers 库中提供；快速词元化器是用 Rust 编写的，可在 Tokenizers 库中提供，可提升性能。如需了解详情，请参阅 Hugging Face 文档。
• --trust_remote_code：是否允许 Hugging Face 模型代码库中定义的远程代码文件。默认值为 False。
• --load_format：要加载的模型检查点的格式。可能的选项为 ["auto", "dummy"]。默认值为 "auto"。如果此参数设置为 "auto"，则以 safetensors 格式加载模型权重。如果此参数设置为 "dummy"，则模型权重会随机初始化。将此值设置为 "dummy" 有助于进行实验。
• --max_model_len：模型可处理的最大上下文长度（输入长度加上输出长度）。默认值从 Hugging Face 格式的模型配置文件 config.json 中读取。最大上下文长度越大，需要的 TPU 内存就越多。
• --sliding_window：如果设置，此参数会替换模型的滑动窗口注意力机制的窗口大小。将此参数设置为较大的值可使注意力机制纳入更多token，并接近标准自注意力机制的效果。此参数仅用于实验性用途。在一般使用情形下，我们建议使用模型的原始窗口大小。
• --seed：用于初始化所有随机数生成器的种子。更改此参数可能会影响针对同一提示生成的输出，因为这会改变作为下一个token进行抽样的token。默认值为 0。

推理引擎

• --num_hosts：要运行的主机数量。默认值为 1。如需了解详情，请参阅有关 TPU v5e 配置的文档。
• --disagg_topo：使用实验性功能“分离式服务”定义预填充工作器和解码工作器的数量。默认值为 None。参数采用以下格式："number_of_prefill_workers,number_of_decode_workers"。
• --data_parallel_size：数据并行副本的数量。默认值为 1。将此值从 1 设置为 N 大致可将吞吐量提高 N，同时保持相同的延迟时间。
• --tensor_parallel_size：张量并行副本的数量。默认值为 1。增加张量并行复制副本的数量通常可以缩短延迟时间，因为这样可以通过减小矩阵大小来加快矩阵乘法运算。
• --worker_distributed_method：用于启动工作器的分布式方法。 对于 multiprocessing 模块，请使用 mp；对于 Ray 库，请使用 ray。默认值为 mp。
• --enable_jit：是否启用 JIT（即时编译）模式。默认值为 True。设置 --no-enable_jit 可将其停用。 启用 JIT 模式可提高推理性能，但需要花费更多时间进行初始编译。一般来说，推理性能的提升会超过开销。
• --warmup：是否在初始化期间通过示例请求预热服务器。默认值为 True。将 --no-warmup 设置为 0 可将其停用。建议进行预热，因为初始请求会触发更重的编译，因此速度会较慢。
• --max_prefill_seqs：每次迭代可安排预填充的序列数上限。默认值为 1。此值越大，服务器可以实现的吞吐量越高，但可能对延迟时间有一定影响。
• --prefill_seqs_padding：服务器将预填充批次大小填充为该值的倍数。默认值为 8。增加此值可缩短模型重新编译时间，但会增加浪费的计算量和推理开销。最佳设置取决于请求流量。
• --prefill_len_padding：服务器将序列长度填充为该值的倍数。默认值为 512。增加此值可缩短模型重新编译时间，但会增加浪费的计算和推理开销。最佳设置取决于请求的数据分布。
• --max_decode_seqs/--max_running_seqs：每次迭代可安排解码的序列数上限。默认值为 256。 此值越大，服务器可以实现的吞吐量越高，但可能对延迟时间有一定影响。
• --decode_seqs_padding：服务器将解码批次大小填充为该值的倍数。默认值为 8。增加此值可缩短模型重新编译时间，但会增加浪费的计算和推理开销。最佳设置取决于请求流量。
• --decode_blocks_padding：在解码期间，服务器会将序列的键值缓存 (KV 缓存) 所用的内存块数量填充为该值的倍数。默认值为 128。增加此值可缩短模型重新编译时间，但会增加浪费的计算和推理开销。最佳设置取决于请求的数据分布。
• --enable_prefix_cache_hbm：是否在 HBM 中启用前缀缓存。默认值为 False。设置此参数可通过重复使用之前请求的共享前缀的计算结果来提高性能。
• --enable_chunked_prefill：是否启用分块预填充。默认值为 False。设置此参数可以支持更长的上下文长度并提高性能。

内存管理

• --hbm_utilization_factor：加载模型权重后，可为 KV 缓存分配的空闲 Cloud TPU 高带宽内存 (HBM) 所占的百分比。默认值为 0.9。将此参数设置为较高的值会增加 KV 缓存大小，并可以提高吞吐量，但会增加在初始化期间和运行时耗尽 Cloud TPU HBM 的风险。
• --num_blocks：要为 KV 缓存分配的设备块数量。如果设置了此参数，服务器会忽略 --hbm_utilization_factor。如果未设置此参数，服务器会分析 HBM 使用情况，并根据 --hbm_utilization_factor 计算要分配的设备块数量。将此参数设置为较高的值会增加 KV 缓存大小，并可以提高吞吐量，但也会增加在初始化期间和运行时耗尽 Cloud TPU HBM 的风险。
• --block_size：存储在一个块中的 token 数量。可能的选项为 [8, 16, 32, 2048, 8192]。默认值为 32。将此参数设置为较大的值可减少块管理中的开销，但会浪费更多内存。确切的性能影响需要通过实证来确定。

动态 LoRA

• --enable_lora：是否启用从 Cloud Storage 动态加载 LoRA 适配器。默认值为 False。Llama 模型系列支持此功能。
• --max_lora_rank：LoRA 适配器支持的 LoRA 秩上限（在请求中定义）。默认值为 16。将此参数设置为较高的值可让服务器更灵活地使用 LoRA 适配器，但会增加为 LoRA 权重分配的 Cloud TPU HBM 量，并降低吞吐量。
• --enable_lora_cache：是否启用动态 LoRA 适配器缓存。 默认值为 True。设置 --no-enable_lora_cache 可将其停用。 缓存可提高性能，因为这样一来，您无需重新下载之前使用过的 LoRA 适配器文件。
• --max_num_mem_cached_lora：存储在 TPU 内存缓存中的 LoRA 适配器数量上限。默认值为 16。将此参数设置为较大的值可提高缓存命中的几率，但会增加 Cloud TPU HBM 用量。

您还可以使用以下环境变量配置服务器：

• HEX_LLM_LOG_LEVEL：控制生成的日志信息量。 默认值为 INFO。将此属性设置为日志记录模块中定义的某个标准 Python 日志记录级别。
• HEX_LLM_VERBOSE_LOG：是否启用详细的日志记录输出。允许的值为 true 或 false。默认值为 false。

调优服务器参数

服务器参数相互关联，共同影响服务性能。例如，如果将 --max_model_len=4096 设置为较大的值，会导致 TPU 内存使用量增加，因此需要分配更多内存并减少批处理。此外，有些参数由使用情形决定，而有些参数可以调优。以下是配置 Hex-LLM 服务器的工作流。

1. 确定感兴趣的模型系列和模型变体。例如，Llama 3.1 8B Instruct。
2. 根据模型大小和精度估算所需 TPU 内存的下限：model_size * (num_bits / 8)。对于 80 亿参数模型和 bfloat16 精度，所需的 TPU 内存下限为 8 * (16 / 8) = 16 GB。
3. 估计所需的 TPU v5e 芯片数量，其中每个 v5e 芯片提供 16 GB 内存：tpu_memory / 16。对于 8B 模型和 bfloat16 精度，您需要使用多个芯片。在单芯片、四芯片和八芯片配置中，提供超过 1 个芯片的最小配置是四芯片配置：ct5lp-hightpu-4t。您可以随后设置 --tensor_parallel_size=4。
4. 确定预期使用情形下的最大上下文长度（输入长度 + 输出长度）。例如，4096。您可以随后设置 --max_model_len=4096。
5. 将为 KV 缓存分配的空闲 TPU 内存量调整为在给定模型、硬件和服务器配置 (--hbm_utilization_factor) 的情况下可实现的最大值。从 0.95 开始。部署 Hex-LLM 服务器，并使用长提示和高并发来测试该服务器。如果服务器内存不足，请相应地降低利用率因子。

以下是用于部署 Llama 3.1 8B Instruct 的一组示例参数：

python -m hex_llm.server.api_server \
    --model=meta-llama/Llama-3.1-8B-Instruct \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.95

在 ct5lp-hightpu-4t 上部署 Llama 3.1 70B Instruct AWQ 的一组示例参数如下：

python -m hex_llm.server.api_server \
    --model=hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4 \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.45

申请 Cloud TPU 配额

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