使用 Hex-LLM 進階容器和 Cloud TPU 提供開放原始碼模型

Hex-LLM 是以 XLA 提供的高效率大型語言模型 (LLM)，也是 Gemini Enterprise Agent Platform LLM 的服務提供框架，經過特別設計與最佳化調整，適用於 Cloud TPU 硬體。Hex-LLM 結合了 LLM 提供技術 (例如持續批次處理和 PagedAttention)，並搭配專為 XLA 和 Cloud TPU 量身打造的 Gemini Enterprise Agent Platform 最佳化功能。這類 LLM 會在 Cloud TPU 提供，可用於開放原始碼模型。

您可以透過模型 Playground、一鍵部署和筆記本，在 Model Garden 中使用 Hex-LLM。

功能與特色

Hex-LLM 以開放原始碼專案為基礎，並針對 XLA 和 Cloud TPU 進行 Google 專屬的最佳化調整。Hex-LLM 在提供常用 LLM 時，可達到高處理量和低延遲。

Hex-LLM 包含下列最佳化功能：

  • 採用以權杖為基礎的連續批次處理演算法，確保模型在處理大量並行要求時，能充分運用硬體資源。
  • 完全重寫注意力核心，並針對 XLA 進行最佳化。
  • 彈性且可組合的資料平行處理和張量平行處理策略，搭配高度最佳化的權重 sharding 方法，可在多個 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，其中最多包含八個 v5e 核心。

如要啟用這項功能，請在 Hex-LLM 容器引數中設定 --num_hosts，並在 Vertex AI SDK 模型上傳要求中設定 --tpu_topology。下列範例說明如何使用 TPU 4x4 v5e 拓撲部署 Hex-LLM 容器，提供 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 容器的端對端教學課程，請參閱 Agent Platform 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"，可設定三個預先填入 worker 和 1 個解碼 worker。每個 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 的倍數。在 Chat Completion 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 Playground 是預先部署的 Agent Platform 端點，只要在模型資訊卡中傳送要求，即可連線。

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

  2. 點選「提交」，即可快速取得模型回覆。

試用 Gemma！

使用一鍵部署

您可以使用模型資訊卡，部署採用 Hex-LLM 的自訂 Agent Platform 端點。

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

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

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

使用 Colab Enterprise 筆記本

如需彈性和自訂功能，您可以使用 Colab Enterprise 筆記本範例，透過 Python 適用的 Agent Platform SDK，部署採用 Hex-LLM 的 Agent Platform 端點。

  1. 前往模型資訊卡頁面，然後按一下「Open 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 筆記本範例包括：

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

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

型號

  • --model：要載入的模型。您可以指定 Hugging Face 模型 ID、Cloud Storage bucket 路徑 (gs://my-bucket/my-model) 或本機路徑。模型構件應採用 Hugging Face 格式，並使用 safetensors 檔案儲存模型權重。Llama、Gemma 2 和 Mistral/Mixtral 支援 BitsAndBytes int8 和 AWQ 量化模型構件。
  • --tokenizer：要載入的權杖化工具。可以是 Hugging Face 模型 ID、Cloud Storage bucket 路徑 (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：定義預先填入 worker 和解碼 worker 的數量，並啟用實驗性功能「分解服務」。預設值為 None。引數格式如下："number_of_prefill_workers,number_of_decode_workers"。
  • --data_parallel_size：資料平行副本數量。預設值為 1。將這個值從 1 設為 N，可將處理量提高約 N，同時維持相同的延遲時間。
  • --tensor_parallel_size：張量平行副本數量。預設值為 1。增加張量平行副本數量，通常可以縮短延遲時間，因為縮減矩陣大小可加快矩陣乘法速度。
  • --worker_distributed_method：啟動 worker 的分散式方法。多重處理模組請使用 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：區塊中儲存的權杖數量。可能選項為 [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)。如果是 8B 模型和 bfloat16 精確度，所需的 TPU 記憶體下限為 8 * (16 / 8) = 16 GB。
  3. 估算所需的 TPU v5e 晶片數量，每個 v5e 晶片提供 16 GB： 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 筆記本部署作業。如要申請提高配額值，請參閱「要求調整配額」。