在 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 中使用 Hex-LLM。

功能

Hex-LLM 以開放原始碼專案為基礎，並採用 Google 針對 XLA 和 Cloud TPU 的最佳化技術。在提供常用的 LLM 時，Hex-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 推理和推理加
• 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 VM (最多含有八個 v5e 核心) 的大型模型。

如要啟用這項功能，請在 Hex-LLM 容器引數中設定 --num_hosts，並在 Vertex AI SDK 模型上傳要求中設定 --tpu_topology。以下範例說明如何部署 Hex-LLM 容器，並搭配提供 Llama 3.1 70B bfloat16 模型的 TPU 4x4 v5e 拓樸：

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 (Deployment) 筆記本。

一般來說，啟用多主機放送功能所需的變更只有：

1. 將引數 --tensor_parallel_size 設為 TPU 拓撲中的核心總數。
2. 將引數 --num_hosts 設為 TPU 拓撲中的主機數量。
3. 使用 Vertex AI SDK 模型上傳 API 設定 --tpu_topology。

不區隔放送 [實驗功能]

Hex-LLM 現在支援分散式服務做為實驗性功能。這項功能只能在單一主機設定中啟用，且效能仍在調整中。

分割式放送是一種有效的方法，可平衡每個要求的首次符記時間 (TTFT) 和輸出符記時間 (TPOT)，以及整體放送傳送量。它會將預填階段和解碼階段分成不同的工作負載，以免彼此干擾。對於設有嚴格延遲時間要求的情況，這個方法特別實用。

如要啟用這項功能，請在 Hex-LLM 容器引數中設定 --disagg_topo。以下範例說明如何在提供 Llama 3.1 8B bfloat16 模型的 TPU v5e-8 上部署 Hex-LLM 容器：

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"，可設定三個預填工作站和 1 個解碼工作站。每個 worker 都會使用兩個 TPU v5e 核心。

前置字串快取

前置字串快取可縮短提示的「首次符記時間」(TTFT)，針對提示開頭有相同內容的情況，例如公司全體適用的前置詞、常見的系統指示，以及多回合對話記錄。Hex-LLM 可保留已處理輸入符記運算的暫時快取，以改善 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 個符記，可使用 prefill_len_padding 進行設定)。快取命中會以這個值的遞增量發生，確保快取符記數一律是 prefill_len_padding 的倍數。在即時通訊完成 API 回應中，usage.prompt_tokens_details 的 cached_tokens 欄位會指出有多少提示符權碼命中快取。

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

分塊預填

分割預先填入會將要求預先填入分割成較小片段，並將預先填入和解碼混合成一個批次步驟。Hex-LLM 會實作區塊預先填充功能，以平衡首次產生符記的時間 (TTFT) 和個別輸出符記的時間 (TPOT)，並改善傳輸量。

如要啟用這項功能，請在 Hex-LLM 容器引數中設定 --enable_chunked_prefill。以下範例說明如何在提供 Llama 3.1 8B 模型的 TPU v5e-8 上部署 Hex-LLM 容器：

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。您可以透過實驗室、一鍵部署，以及各種模型的 Colab Enterprise 筆記本範例，存取這項技術。

使用遊樂場

Model Garden 遊樂場是預先部署的 Vertex AI 端點，可透過在模型資訊卡中傳送要求來存取。

1. 輸入提示，並視需要加入要求的引數。

2. 按一下「提交」即可快速取得模型回覆。

試試看 Gemma！

使用一鍵部署

您可以使用模型資訊卡，透過 Hex-LLM 部署自訂的 Vertex AI 端點。

1. 前往模型資訊卡頁面，然後按一下「部署」。

2. 針對要使用的模型變化版本，請選取Cloud TPU v5e 機器類型進行部署。

3. 按一下底部的「部署」，即可開始部署程序。您會收到兩封電子郵件通知，一封是在模型上傳時，另一封是在端點就緒時。

使用 Colab Enterprise 筆記本

為求靈活性和客製化，您可以使用 Colab Enterprise 筆記本範例，透過 Python 適用的 Vertex AI SDK 部署採用 Hex-LLM 的 Vertex AI 端點。

1. 前往模型資訊卡頁面，然後按一下「開啟 Notebook」。

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 檔案的模型權重。BitsAndBytes 支援 Llama、Gemma 2 和 Mistral/Mixtral 的 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：如果已設定，這個引數會覆寫模型的滑動視窗注意力視窗大小。將這個引數設為較大的值，可讓注意力機制納入更多符記，並接近標準自我注意力機制的效果。這個引數僅供實驗使用。在一般用途中，建議使用模型的原始視窗大小。
• --seed：用於初始化所有隨機號碼產生器的種子。變更這個引數可能會影響同一個提示產生的輸出內容，因為這會變更用作下一個符號的符號。預設值為 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：啟動 worker 的 distributed 方法。請為多重處理模組使用 mp，或為 Ray 程式庫使用 ray。預設值為 mp。
• --enable_jit：是否要啟用 JIT (即時編譯) 模式。預設值為 True。將 --no-enable_jit 設為 則可停用。啟用 JIT 模式可改善推論效能，但需要額外花費時間進行初始編譯。一般來說，推論效能的好處會大於額外負擔。
• --warmup：在初始化期間，是否要使用範例要求暖機伺服器。預設值為 True。將 --no-warmup 設為停用。建議您進行暖機，因為初始要求會觸發較重的編譯作業，因此速度會變慢。
• --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：儲存在區塊中的符記數量。可能的選項為 [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。將此值設為 logging 模組 中定義的標準 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)。對於 8B 模型和 bfloat16 精確度，所需的 TPU 記憶體下限為 8 * (16 / 8) = 16 GB。
3. 估算所需的 TPU v5e 晶片數量，每個 v5e 晶片提供 16GB： tpu_memory / 16。如要使用 8B 模型和 bfloat16 精度，您需要使用多個晶片。在 1 個晶片、4 個晶片和 8 個晶片設定中，提供超過 1 個晶片的最小設定為 4 個晶片設定：ct5lp-hightpu-4t。接著，您可以設定 --tensor_parallel_size=4。
4. 判斷所需用途的背景資訊長度上限 (輸入長度 + 輸出長度)。例如 4096。接著，您可以設定 --max_model_len=4096。
5. 根據模型、硬體和伺服器設定 (--hbm_utilization_factor)，將為 KV 快取分配的空閒 TPU 記憶體量調校至可達到的最大值。請從 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 筆記本部署作業。如要申請提高配額值，請參閱「申請配額調整」。