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) 筆記本。
一般來說，啟用多主機放送功能只需要進行下列變更：
- 將
--tensor_parallel_size引數設為 TPU 拓撲中的核心總數。 - 將引數
--num_hosts設為 TPU 拓撲中的主機數量。 - 使用 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 端點，只要在模型資訊卡中傳送要求，即可存取。
輸入提示，並視需要加入要求的引數。
按一下「提交」，即可快速取得模型回覆。
使用一鍵部署功能
您可以使用模型資訊卡，透過 Hex-LLM 部署自訂 Vertex AI 端點。
前往模型資訊卡頁面，然後按一下「部署」。
如要使用模型變體，請選取 Cloud TPU v5e 機器類型進行部署。
按一下底部的「部署」，即可開始部署程序。您會收到兩封電子郵件通知，一封是在模型上傳時收到，另一封則是在端點準備就緒時收到。
使用 Colab Enterprise 筆記本
如要享有彈性和自訂功能，可以使用 Colab Enterprise 筆記本範例，透過 Vertex AI SDK for Python 部署搭載 Hex-LLM 的 Vertex AI 端點。
前往模型資訊卡頁面，然後按一下「開啟筆記本」。
選取 Vertex Serving 筆記本。筆記本會在 Colab Enterprise 中開啟。
執行筆記本，使用 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 伺服器的工作流程。
- 判斷感興趣的模型系列和模型變體。例如 Llama 3.1 8B Instruct。
- 根據模型大小和精確度估算所需的 TPU 記憶體下限：
model_size * (num_bits / 8)。如果是 8B 模型和 bfloat16 精確度，TPU 記憶體下限為8 * (16 / 8) = 16 GB。 - 估算所需的 TPU v5e 晶片數量，每個 v5e 晶片提供 16 GB：
tpu_memory / 16。如果是 8B 模型和 bfloat16 精度，您需要超過 1 個晶片。在單晶片、四晶片和八晶片設定中，提供超過 1 個晶片的最小設定是四晶片設定：ct5lp-hightpu-4t。您之後可以設定--tensor_parallel_size=4。 - 根據預期用途，判斷背景資訊長度上限 (輸入長度 + 輸出長度)。例如 4096。您之後可以設定
--max_model_len=4096。 - 將分配給 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 筆記本部署作業。如要申請提高配額值，請參閱「申請調整配額」一文。