使用 TPU 自動調度大型語言模型 (LLM) 推論工作負載的最佳做法


本最佳做法指南說明可用的指標,以及如何選取合適的指標,在 GKE 上為單一主機 JetStream 推論工作負載設定 水平 Pod 自動調度器(HPA),並搭配 TPU 使用。HPA 可確保模型伺服器隨著負載適當擴充,是相當有效率的做法。微調 HPA 設定是主要方法,可根據流量需求調整佈建的硬體成本,以達成推論伺服器效能目標。

如需如何實作這些最佳做法的範例,請參閱「使用 GKE 在 TPU 上為 LLM 工作負載設定自動調度資源」。

目標

本指南適用於生成式 AI 客戶、新舊 GKE 使用者、機器學習工程師,以及有興趣使用 Kubernetes 和 TPU 最佳化單一主機 JetStream 工作負載的 LLMOps (DevOps) 工程師。

閱讀本指南後,您應該能夠:

  • 瞭解單一主機 JetStream 推論的自動調度資源主要指標。
  • 瞭解考慮要自動調度哪些指標時,高層級的取捨。

JetStream 推論作業的自動調度資源指標總覽

建議您使用下列指標:

伺服器指標

JetStream 與許多其他 LLM 推論伺服器一樣,會提供效能指標。GKE 會根據這些伺服器層級指標,簡化 JetStream 的監控和自動調整資源配置作業。如要設定任何自動調度資源功能,您必須先瞭解 JetStreams 請求管道如何影響重要效能指標。所有傳入要求都會經過預先填入佇列、轉移佇列和產生佇列,然後成為解碼要求。JetStream 會持續接受解碼要求,並使用固定數量的解碼執行緒同時處理這些要求。在特定時間點,用於處理解碼要求的解碼引擎所占百分比會以 jetstream_slots_used_percentage 指標表示。

如要擴充單一主機 JetStream,這會對延遲時間和處理量造成以下兩項影響:

  • 如果傳入請求的速率低於解碼時段處理請求的速率,佇列就不會積壓請求。如果 JetStream 沒有待處理事項,輸送量就會與傳入要求的速率成正比。延遲時間大致會維持不變,但會略為增加,且與並行解碼要求數量成正比,因為新接受的解碼要求會中斷解碼作業。
  • 一旦傳入要求的速度超過解碼位置處理要求的速度,要求就會在佇列中積壓。如果 JetStream 有待處理要求,要求延遲時間會大幅增加,且與待處理要求數量成正比,但總處理量會維持不變。

下列伺服器指標與這些相關成效指標的關聯性最強,建議您使用這些指標進行自動調度:

這些指標通常不受效能和流量波動影響,因此是各種 TPU 硬體設定自動調度的可靠起點。

TPU 指標

由於 LLM 服務受記憶體限制,而非運算資源,因此建議您根據記憶體用量調整 JetStream 規模,而非其他 TPU 指標,因為記憶體用量最能反映硬體的資源使用率。詳情請參閱相關的最佳做法一節。

選擇自動調度資源指標時的考量事項

請參考下列注意事項和最佳做法,在 GKE 上選取最適合自動調度資源的指標,以達成推論工作負載的效能目標。

最佳做法:使用預先填入的待處理項目 (佇列) 大小,在特定目標延遲時間門檻內,盡量提高總處理量並降低成本

如果想盡量提高輸送量並降低成本,且模型伺服器每個裝置批次大小的最大輸送量可達到延遲時間目標,建議預先填入佇列大小自動調度資源。

預先填入佇列大小與要求延遲時間直接相關。傳入的要求會在模型伺服器的預先填入佇列中排隊,然後才會處理,而這段佇列時間會增加整體延遲時間。佇列大小是負載尖峰的敏感指標,因為負載增加會迅速填滿佇列。

根據預先填入佇列大小自動調整資源配置,可在負載增加時擴充資源,並在佇列為空時縮減資源,藉此縮短佇列時間。這個方法相對容易實作,因為佇列大小與要求大小、模型或硬體無關。

如要盡量提高每個模型伺服器副本的總處理量,請考慮著重於預填佇列大小。預先填入佇列大小會追蹤待處理的要求,而非處理中的要求。JetStream 使用連續批次處理,可在批次空間可用時,盡量處理並行要求,並維持低佇列。批次空間有限時,佇列會明顯變長,因此請將成長點做為啟動擴充的信號。結合佇列大小自動調整功能和最佳化批次總處理量,即可盡量提高要求總處理量。

判斷 HPA 的最佳佇列大小門檻值

如要選擇正確的佇列大小門檻,請從 3 到 5 之間的值開始,然後逐步增加,直到要求達到偏好的延遲時間為止。使用 locust-load-inference 工具進行測試。如果閾值低於 10,請微調 HPA 擴充設定,以處理流量尖峰。

您也可以建立 Cloud Monitoring 自訂資訊主頁,以視覺化方式呈現指標行為。

限制

請注意 HPA 容許度,預設值為目標值周圍 0.1 的無動作範圍,可抑制震盪。

預先填入佇列大小不會直接控管並行要求,因此其門檻無法保證延遲時間低於每個裝置的批次大小。解決方法是手動減少每個裝置的批次大小,或根據批次大小自動調整資源配置

最佳做法:使用 slots_used 百分比,達到比佇列大小更低的目標延遲時間門檻

如果您的工作負載對延遲時間很敏感,且以佇列為準的調度資源速度不夠快,無法滿足需求,建議選擇以 slots_used 為準的自動調度資源。

自動調度運算單元 (slots_used) 可確保工作負載輸送量擴充,盡可能同時處理大量要求,並在同時處理的要求較少時縮減規模。這會對延遲時間造成兩項影響。首先,由於系統會根據 slots_used 進行自動調度,確保每個傳入要求都有可用的時段,因此門檻越接近 1,要求排隊等待的時間就越長,延遲時間也就越長。其次,批次越大,總處理量就越高,但由於部分要求的預先填入階段會中斷連續批次模型伺服器中其他要求的解碼階段,因此延遲時間也會增加。您可以監控批次大小模式,並使用自動調度資源功能,在負載尖峰期間盡量減少並行要求。

如果佇列大小已符合延遲目標,請優先進行自動調度資源。這可大幅提高輸送量和成本效益。不過,slots_used 對於容易受到延遲影響的工作負載來說很有價值。

此外,建議您先將每個裝置的批次大小調整為適當值,再探索以 slots_used 為依據的自動調度功能。您也可以選擇搭配使用以佇列為準的自動調度資源功能。

為 HPA 決定最佳 slots_used 百分比門檻值

如要選擇合適的批次大小門檻,請實驗性地增加伺服器負載,並觀察批次大小達到峰值的位置。此外,我們建議使用 locust-load-inference 工具進行測試。找出可將記憶體用量最大化的最佳單一裝置批次大小後,即可設定 slots_used 百分比目標。將初始目標值設為略低於 1,然後調低目標值,直到達到偏好的延遲時間為止。

您也可以建立 Cloud Monitoring 自訂資訊主頁,以視覺化方式呈現指標行為。

限制

請注意 HPA 容許值,這是目標值周圍的預設 0.1 無動作範圍,可抑制震盪。

雖然根據 slots_used 百分比自動調整資源有助於控制延遲,但仍有其限制。由於要求大小和硬體限制各不相同,因此很難找出適當的 slots_used 百分比門檻。如果縮放規則嘗試將平均 slots_used 百分比維持在 100% 以下,表示縮放規則嘗試保留非零數量的可用時段。這些可用時段對應到未使用的可用輸送量,如果您想充分運用可用的 TPU,這並非理想做法。

最佳做法:使用 TPU 高頻寬記憶體 (HBM) 提高硬體使用率

即使與系統指標相比,TPU 高頻寬記憶體 (HBM) 用量與硬體使用率的對應關係也最直接,因為系統指標不會將要求大小納入考量。雖然使用佇列大小進行調整可盡量提高硬體使用率,但延遲時間也會增加。如果您偏好使用系統指標而非伺服器指標,建議您使用 HBM 用量,這是依據 slots_used 自動調度資源的最佳替代方案,因為兩者都對應到輸送量。如要進一步瞭解 TPU 記憶體,請參閱「TPU 運作方式」。

如果批次大小超過最佳值,總處理量會提高,但記憶體不足 (OOM) 錯誤的風險也會增加。我們強烈建議您根據 HBM 指標進行調整,以平衡輸送量和穩定性。建議不要同時根據預填佇列大小和 HBM 使用量進行調整,因為負載增加時,HBM 使用量會增加,並先觸發調整。

為 HPA 決定最佳 TPU HBM 使用率門檻值

建議您先將每個裝置的批次大小設為某個值,在預期負載量達到上限時,盡可能提高記憶體用量,再選擇記憶體用量門檻。請注意,這個值必須透過實驗判斷,且很大程度取決於 HBM 總量,以及預期的提示和回覆長度。建議您使用 locust-load-inference 工具進行實驗和測試。

與每個裝置的批次大小類似,請設定門檻,盡量提高 TPU 記憶體使用率,同時將 OOM 行為的風險降到最低。

您也可以建立 Cloud Monitoring 自訂資訊主頁,以視覺化方式呈現指標行為。

限制

有兩項注意事項會削弱延遲時間和輸送量與 HBM 使用量之間的對應關係,也就是 HBM 使用量的波動性,以及一般 TPU 指標的取樣率較低。此外,雖然 HBM 使用量與延遲時間之間仍存在對應關係,但 HBM 使用量增加對延遲時間的影響,遠小於佇列要求數量增加的影響。

最佳做法:最佳化 HPA 設定

建議您設定下列 HPA 設定選項:

  • 穩定期:使用這個 HPA 設定選項,防止指標波動導致副本數量快速變更。預設值為向下擴充 5 分鐘 (避免過早向下擴充),向上擴充則為 0 分鐘 (確保回應能力)。根據工作負載的波動和偏好的回應速度調整值。
  • 資源調度政策:使用這個 HPA 設定選項,微調資源調度行為。您可以設定「Pod」政策限制,指定每個時間單位變更的絕對副本數,並設定「百分比」政策限制,指定變更的百分比。

如要進一步瞭解這些選項,請參閱開放原始碼 Kubernetes 說明文件的「水平 Pod 自動調度」一節。

後續步驟