在 Cloud TPU 上使用 Hex-LLM 高級容器提供開放模型

Hex-LLM 是以 XLA 提供的高效率、低成本大型語言模型 (LLM)，也是專為 Cloud TPU 硬體設計與最佳化的 Vertex AI LLM 提供框架。Hex-LLM 結合了 LLM 提供技術 (例如持續批次處理和 PagedAttention)，並搭配專為 XLA 和 Cloud TPU 量身打造的 Vertex AI 最佳化功能。這類 LLM 會在 Cloud TPU 提供，可用於開放原始碼模型。

您可以在模型花園中透過模型遊樂場、一鍵部署和筆記本使用 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 reasoning 和 reasoning 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 bucket 載入 Hugging Face 模型。
  • 使用 bitsandbytes 和 AWQ 進行量化。
  • 動態載入 LoRA。Hex-LLM 能夠在服務期間讀取要求引數，藉此載入 LoRA 權重。

進階功能

Hex-LLM 支援下列進階功能：

  • 多主機放送
  • 放送不相關的廣告 [實驗功能]
  • 前置字串快取
  • 支援 4 位元量化

多主機放送

Hex-LLM 現在支援透過多主機 TPU 切片提供模型。這項功能可讓您提供無法載入單一主機 TPU VM 的大型模型，這類 VM 最多包含八個 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,
)

如需部署 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。 以下範例說明如何在 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"，可設定三個預先填入工作站和 1 個解碼工作站。每個工作站會使用兩個 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。以下範例說明如何在 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 試用版是預先部署的 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 筆記本範例包括：

設定伺服器引數和環境變數

您可以設定下列引數來啟動 Hex-LLM 伺服器。您可以調整引數，以最符合預期用途和需求。請注意，這些引數是預先定義的，可供一鍵部署，提供最簡單的部署體驗。如要自訂引數，可以參考筆記本範例，並據此設定引數。

型號

  • --model：要載入的模型。您可以指定 Hugging Face 模型 ID、Cloud Storage bucket 路徑 (gs://my-bucket/my-model) 或本機路徑。模型構件應採用 Hugging Face 格式，並使用 safetensors 檔案儲存模型權重。BitsAndBytes int8 和 AWQ 量化模型構件支援 Llama、Gemma 2 和 Mistral/Mixtral。
  • --tokenizer：要載入的 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：啟動工作者的分散式方法。 請使用 mp 搭配 multiprocessing 模組，或使用 ray 搭配 Ray 程式庫。預設值為 mp。
  • --enable_jit：是否啟用 JIT (Just-in-Time Compilation) 模式。預設值為 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 晶片提供 16 GB： tpu_memory / 16。如果是 8B 模型和 bfloat16 精度，您需要超過 1 個晶片。在單晶片、四晶片和八晶片設定中，提供超過 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 筆記本部署作業。如要申請提高配額值，請參閱「申請調整配額」一文。