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 筆記本部署作業。如要申請提高配額值,請參閱「申請調整配額」一文。