SageMaker HyperPod、モデルキャッシュで推論のコールドスタートを短縮

何が発表されたのか

AWSは、Amazon SageMaker HyperPod上の推論(Amazon SageMaker Inference on HyperPod)向けに「モデルキャッシュ」を発表しました。推論ポッドが起動してトラフィックを処理できるようになるまでの待ち時間、いわゆるコールドスタートを短縮する機能です。HyperPodはKubernetes上で大規模なモデル学習・推論を動かすための仕組みで、ポッドの起動時にコンテナイメージとモデルの重みを毎回ダウンロードする必要がありました。

この待ち時間は2つのダウンロードに支配されています。1つはAmazon Elastic Container Registry(ECR)からの推論サーバーのコンテナイメージ、もう1つはAmazon S3、Amazon FSx for Lustre、HuggingFace Hubなどから取得するモデルの重みです。原文によれば、vLLMやLMIといった推論サーバーのイメージは数GB規模で取得に5〜7分、145GBのモデルをS3から取得する場合はさらに20分以上かかるといいます。

600GBを超えるDeepSeek-R1のような大規模モデルでは、1件のリクエストを処理するまでに30分以上かかると説明されています。モデルキャッシュを有効にすると、ポッドはローカルのNVMeストレージから約7GB/sで読み出せるようになり、数秒でトラフィックを処理し始められるとしています。

2つのキャッシュ機能

モデルキャッシュは「重みキャッシュ」と「イメージキャッシュ」という独立した2つの機能で構成されます。どちらか一方だけを有効にすることも、両方を同時に有効にすることもできます。設定は既存のInferenceEndpointConfigまたはJumpStartModelリソースにmodelCacheConfigセクションを追加するだけで、追加のインフラ構築は不要とされています。

重みキャッシュでは、HyperPod Inference OperatorがModelDataCacheConfigリソースを自動作成し、指定したソースから対象ノードのローカルNVMeへ重みを事前ダウンロードします。ダウンロードが完了したノードにはcache-readyというラベルが付き、すべての対象ノードがcache-readyになるまで推論デプロイは作成されません。キャッシュは同じノード上でのポッド再起動をまたいで保持されるとしています。

イメージキャッシュでは、OperatorがDaemonSetを作成して対象ノードすべてにコンテナイメージを事前取得します。重みキャッシュと異なり、デプロイの作成はブロックしません。キャッシュ済みのノードではECRからの取得を完全にスキップでき、5〜7分の短縮につながります。同じイメージを使う複数のデプロイは1つのキャッシュリソースを共有し、参照がなくなった時点でOperatorが削除します。

フォールバックとCRD管理

両機能ともスケジューリングは「必須」ではなく「優先」で行われます。ポッドはキャッシュ済みノードを優先しますが、起動をブロックされることはありません。急激なスケールアウトでキャッシュ済みノード数を超えた場合などは、従来どおりS3やFSxから重みを読み込み、ECRからイメージを取得します。原文はこれを、障害でもユーザー介入でもなく、通常のダウンロード時間を超える劣化もないと説明しています。

キャッシュのライフサイクルは2つのCustom Resource Definition(CRD)で管理されます。Operatorが自動作成・管理するため、利用者が直接作成する必要はありません。kubectl get modeldatacacheconfigやkubectl get inferenceimagecacheで状態を確認でき、原文ではSTATEやPHASE、TARGET、READYといった列を持つ出力例が示されています。

モデルの取得元を変更した場合、たとえば新しいS3パスを指定した場合は、Operatorが新しいキャッシュを作成し、更新後のデプロイをロールアウトしてから古いキャッシュを削除します。イメージ変更も同様で、古いデータが残らないゼロダウンタイムの切り替えになるとしています。

ベンチマークと制約

原文によれば、57〜145GBのモデルでは重みキャッシュによりスケールアウトが約60%高速化しました。イメージキャッシュはコールド時のイメージ取得から2分以上を削減し、毎回ECRから取得する場合と比べて最大97%の削減を達成したとしています。モデルが大きいほど、ネットワーク経由でダウンロードしていたデータ量が多いため効果は大きくなり、600GB超のDeepSeek-R1では30分以上かかっていたダウンロードがなくなるといいます。

一方で注意点も示されています。重みキャッシュはノード単位で保持されるため、NVMeの消費量はノード数に比例して増えます。初回のキャッシュ作成時にはリモートソースからのダウンロードが必要で、そのコストは一度だけ発生します。NVMeの容量は有限で、たとえば300GBのモデルに対してNVMeが250GBしかないインスタンスタイプではキャッシュできません。

取得元の更新は自動検出されない点も明記されています。同じS3パスでモデルファイルを更新しても、InferenceEndpointConfigのspecを変更しなければOperatorはキャッシュ済みのバージョンを提供し続けます。新しい重みを反映するには、モデルパスの変更やバージョン接尾辞の追加などspecの更新が必要です。