このベスト プラクティス ガイドでは、利用可能な指標と、GKE 上の TPU を使用した単一ホストの JetStream 推論ワークロードに HorizontalPodAutoscaler（HPA）を設定する際に適切な指標を選択する方法について説明します。HPA は、モデルサーバーが負荷に応じて適切にスケーリングされるようにする効率的な方法です。プロビジョニングされたハードウェアの費用をトラフィックの需要に合わせて調整し、推論サーバーのパフォーマンス目標を達成するには、HPA 設定を微調整するのが主な方法です。
これらのベスト プラクティスを実装する方法の例については、GKE で TPU での LLM ワークロードの自動スケーリングを構成するをご覧ください。
目標
このガイドは、生成 AI をご利用のお客様、GKE の新規または既存のユーザー、ML エンジニア、LLMOps（DevOps）エンジニアで、Kubernetes で TPU を使用して単一ホストの JetStream ワークロードを最適化することに関心のある方を対象としています。
このガイドを読むことで、次のことが可能になります。
- 単一ホストの JetStream 推論の主要な自動スケーリング指標について理解する。
- どの指標で自動スケーリングするかを検討する際のトレードオフを理解する。
JetStream 推論の自動スケーリング指標の概要
次の指標を使用することをおすすめします。
サーバー指標
JetStream は、他の多くの LLM 推論サーバーと同様に、パフォーマンス指標を提供します。GKE は、これらのサーバーレベルの指標に基づいて JetStream のモニタリングと自動スケーリングを簡素化します。自動スケーリングを構成するには、まず、JetStreams リクエスト パイプラインが主要なパフォーマンス指標にどのように影響するかを理解する必要があります。受信したすべてのリクエストは、プリフィル キュー、転送キュー、生成キューを経由してデコード リクエストになります。JetStream は、デコード リクエストをローリング方式で受信し、固定数のデコード スレッドを使用して同時に処理します。特定の時点でデコード リクエストの処理に使用されているデコード エンジンの割合は、jetstream_slots_used_percentage 指標として測定されます。
単一ホストの JetStream をスケーリングする場合、レイテンシとスループットに次の 2 つの影響があります。
- 受信するリクエストのレートが、デコード スロットで処理可能なレートを下回る場合、リクエストはキューにバックログされません。JetStream にバックログがない場合、スループットは受信リクエストのレートに比例します。レイテンシはほとんど一定ですが、新たに受け入れられたデコード リクエストによってデコードが中断されるため、同時実行デコード リクエストの数に比例してわずかに増加します。
- 受信するリクエストのレートが、デコード スロットで処理可能なレートを超えると、リクエストはキューにバックログされます。JetStream にバックログがある場合、スループットは一定のままで、リクエスト レイテンシはバックログされたリクエスト数に比例して大幅に増加します。
次のサーバー指標は、関連するパフォーマンス指標と相関性が最も高いため、自動スケーリングに使用することをおすすめします。
jetstream_prefill_backlog_size: サーバーキュー（バックログ）で処理を待機しているリクエストの数。この指標はレイテンシと強い相関関係があります。詳細については、関連するベスト プラクティスのセクションをご覧ください。jetstream_slots_used_percentage: 推論処理中のリクエスト数（JetStream が処理できるリクエストの合計数の割合）。この指標はスループットと強い相関関係があります。詳細については、関連するベスト プラクティスのセクションをご覧ください。
これらの指標は、パフォーマンスとトラフィックの変動に強いため、さまざまな TPU ハードウェア設定で自動スケーリングを行う際に信頼できる出発点となります。
TPU 指標
LLM のサービングはコンピューティングではなく、メモリによってボトルネックになるため、ハードウェアのリソース使用率を最もよく反映する TPU の指標ではなく、メモリ使用量で JetStream をスケーリングすることをおすすめします。詳細については、関連するベスト プラクティスのセクションをご覧ください。
自動スケーリングの指標を選択する際の考慮事項
次の考慮事項とベスト プラクティスを使用して、推論ワークロードのパフォーマンス目標を達成するために GKE で自動スケーリングに最適な指標を選択します。
ベスト プラクティス: プリフィル バックログ（キュー）サイズを使用して、特定のターゲット レイテンシしきい値内でスループットを最大化し、費用を最小限に抑える
スループットと費用を最適化する場合、また、デバイスごとのバッチサイズに対するモデルサーバーの最大スループットでレイテンシ目標を達成できる場合は、プリフィル キューのサイズによる自動スケーリングをおすすめします。
プリフィル キューのサイズは、リクエストのレイテンシと直接的に関連しています。受信リクエストは、処理される前にモデルサーバーのプリフィル キューに追加されます。このキュー時間は全体的なレイテンシに追加されます。キューのサイズは負荷の急増を示す敏感な指標です。負荷が増加すると、キューがすぐにいっぱいになります。
プリフィル キューのサイズに基づく自動スケーリングでは、負荷が高いときにスケールアップし、キューが空になったときにスケールダウンすることで、キュー時間を最小限に抑えます。キューサイズがリクエスト サイズ、モデル、ハードウェアとは無関係であるため、このアプローチは比較的簡単に実装できます。
各モデルサーバーのレプリカのスループットを最大化するには、プリフィル キューのサイズに注目することを検討してください。プリフィル キューのサイズは、処理中ではなく保留中のリクエストを追跡します。JetStream は連続バッチ処理を使用します。これにより、同時リクエストを最大化し、バッチスペースが使用可能なときのキューを短くすることができます。バッチスペースが限られている場合、キューは著しく増加するため、増加ポイントをスケールアップの開始シグナルとして使用します。キューサイズの自動スケーリングと最適化されたバッチ スループットを組み合わせることで、リクエスト スループットを最大化できます。
HPA に最適なキューサイズのしきい値を決定する
適切なキューサイズのしきい値を選択するには、最初は 3～5 の値から始め、リクエストが目的のレイテンシに達するまで、この値を徐々に増やしていきます。テストには locust-load-inference ツールを使用します。しきい値が 10 未満の場合は、トラフィックの急増に対応するように HPA スケールアップ設定を微調整します。
Cloud Monitoring カスタム ダッシュボードを作成して、指標の動作を可視化することもできます。
制限事項
HPA の許容範囲に注意してください。デフォルトでは、減衰振動のため目標値の前後に 0.1 の無動作範囲が設定されています。
プリフィル キューのサイズは、同時実行リクエストを直接制御しません。しきい値を使用しても、デバイスごとのバッチサイズで許容される値よりもレイテンシが低くなるとは限りません。回避策として、デバイスごとのバッチサイズを手動で減らすか、バッチサイズで自動スケーリングします。
ベスト プラクティス: slots_used の割合を使用して、キューサイズよりも低いターゲット レイテンシしきい値に到達させる
キューベースのスケーリングでは要件を満たせない、レイテンシの影響を受けやすいワークロードがある場合は、slots_used ベースの自動スケーリングを選択することをおすすめします。
slots_used で自動スケーリングを行うと、ワークロードのスループットがスケールアップされ、一度に並列で処理されるリクエスト数が最大化されます。並列で処理されるリクエスト数が少ない場合は、スケールダウンされます。これにより、レイテンシに 2 つの影響があります。まず、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）の使用量は、リクエスト サイズを考慮していないため、システム指標と比較しても、ハードウェア使用率と最も直接的に対応しています。キューサイズでスケーリングすると、ハードウェア使用率を最大化できますが、レイテンシが犠牲になります。サーバー指標ではなくシステム指標を使用する場合は、どちらもスループットに関連するため、slots_used による自動スケーリングではなく HBM 使用量を使用することをおすすめします。TPU メモリについて詳しくは、TPU の仕組みをご覧ください。
バッチサイズを最適なサイズより大きくすると、スループットは向上しますが、メモリ不足（OOM）エラーが発生するリスクも高まります。スループットと安定性のバランスを取るために、HBM 指標に基づいてスケーリングすることを強くおすすめします。負荷が増加すると HBM の使用量が増加し、最初にスケーリングがトリガーされるため、プリフィル キューのサイズと HBM の使用量を同時にスケーリングすることはおすすめしません。
HPA の最適な TPU HBM 使用率しきい値を決定する
メモリ使用率のしきい値を選択する前に、想定される最大負荷で動作する際に使用されるメモリを最大限にするように、デバイスごとのバッチサイズを設定することをおすすめします。この値はテストで決定する必要があります。これは、HBM の合計だけでなく、想定されるプロンプトとレスポンスの長さにも大きく依存します。テストと検証には、locust-load-inference ツールを使用することをおすすめします。
デバイスごとのバッチサイズと同様に、OOM 動作のリスクを最小限に抑えながら TPU メモリ使用率を最大化するようにしきい値を設定します。
Cloud Monitoring カスタム ダッシュボードを作成して、指標の動作を可視化することもできます。
制限事項
レイテンシとスループットは HBM の使用量に対応していますが、この利点を損なう 2 つの注意点があります。一つは、HBM の使用量が不安定であること、もう一つは一般的な TPU 指標のサンプリング レートが低いことです。また、HBM の使用量とレイテンシの間には相関関係がありますが、HBM の使用量の増加がレイテンシに与える影響は、キューに入れられたリクエスト数の増加に比べると、かなり小さくなります。
ベスト プラクティス: HPA 構成を最適化する
次の HPA 構成オプションを設定することをおすすめします。
- 安定化ウィンドウ: この HPA 構成オプションを使用すると、指標の変動によるレプリカ数の急激な変化を防ぐことができます。デフォルトは、スケールダウンの場合は 5 分（早期のスケーリングを回避）、スケールアップの場合は 0（応答性を保証）です。ワークロードの不安定性と優先する応答性に応じて値を調整します。
- スケーリング ポリシー: この HPA 構成オプションを使用して、スケールアップとスケールダウンの動作を微調整します。Pods ポリシーの制限を設定して、時間単位で変更されるレプリカの絶対数を指定できます。また、Percent ポリシーの制限を設定して、変化率を指定できます。
これらのオプションの詳細については、オープンソースの Kubernetes ドキュメントの水平 Pod 自動スケーリングをご覧ください。
次のステップ
- TPU 上の LLM ワークロードの自動スケーリングを構成する方法を学習する。