Run:ai Model Streamer と vLLM を使用して GKE でのモデルの読み込みを高速化する

Autopilot Standard

このドキュメントでは、Google Kubernetes Engine（GKE）で vLLM 推論サーバーとともに Run:ai Model Streamer を使用して、Cloud Storage から大規模な AI モデルの重みを読み込む速度を上げる方法について説明します。

このドキュメントのソリューションでは、safetensors 形式の AI モデルと重みが Cloud Storage バケットにすでに読み込まれていることを前提としています。

vLLM デプロイに --load-format=runai_streamer フラグを追加すると、Run:ai Model Streamer を使用して、GKE 上の AI ワークロードのモデルダウンロード効率を向上させることができます。

このドキュメントは、次のユーザーを対象としています。

オブジェクトストレージから GPU / TPU ノードに大規模な AI モデルをできるだけ早く読み込む必要がある ML エンジニア。
GKE でモデル提供インフラストラクチャを自動化して最適化するプラットフォーム管理者とオペレーター。
AI / ML ワークロード用の専用データ読み込みツールを評価するクラウドアーキテクト。

Google Cloud のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な GKE ユーザーロールとタスクをご覧ください。

概要

このドキュメントで説明するソリューションは、Run:ai Model Streamer、vLLM、safetensors ファイル形式の 3 つのコアコンポーネントを使用して、Cloud Storage から GPU ノードへのモデルの重みの読み込みプロセスを高速化します。

Run:ai Model Streamer

Run:ai Model Streamer は、大規模な AI モデルを GPU に読み込む速度を向上させるオープンソースの Python SDK です。モデルの重みを Cloud Storage バケットなどのストレージから GPU のメモリに直接ストリーミングします。モデルストリーマーは、Cloud Storage にある safetensors ファイルへのアクセスに特に適しています。

safetensors

safetensors は、AI モデルのコアデータ構造であるテンソルを、セキュリティと速度の両方を強化する方法で保存するファイル形式です。safetensors は Python の pickle 形式の代替として設計されており、ゼロコピーアプローチにより読み込み時間を短縮できます。このアプローチでは、ファイル全体をローカルメモリに読み込む必要がなく、ソースからテンソルに直接アクセスできます。

vLLM

vLLM は、LLM の推論およびサービングを行うためのオープンソースライブラリです。大規模な AI モデルを高速に読み込むように最適化された高性能の推論サーバーです。このドキュメントでは、vLLM は GKE で AI モデルを実行し、受信した推論リクエストを処理するコアエンジンとして機能します。Run:ai Model Streamer の Cloud Storage の組み込み認証サポートには、vLLM バージョン 0.11.1 以降が必要です。

Run:ai Model Streamer がモデルの読み込みを高速化する仕組み

推論のために LLM ベースの AI アプリケーションを起動すると、モデルを使用できるようになるまでに大幅な遅延が発生することがよくあります。この遅延はコールドスタートと呼ばれ、数ギガバイトのモデルファイル全体を Cloud Storage バケットなどのストレージの場所からマシンのローカルディスクにダウンロードする必要があるために発生します。ファイルは GPU のメモリに読み込まれます。読み込みの間、高価な GPU がアイドル状態になるため、非効率的でコストがかかります。

このモデルストリーマーは、ダウンロードしてから読み込むプロセスではなく、Cloud Storage から GPU のメモリにモデルを直接ストリーミングします。そして、高性能バックエンドを使用して、モデルの複数の要素（テンソル）を並列で読み取ります。テンソルを同時に読み取る方が、ファイルを順番に読み込むよりも大幅に高速です。

アーキテクチャの概要

Run:ai Model Streamer は GKE の vLLM と統合され、モデルの重みを Cloud Storage から GPU メモリに直接ストリーミングしてローカルディスクをバイパスすることで、モデルの読み込みを高速化します。

次の図は、このアーキテクチャを示しています。

Run:ai Model Streamer が Cloud Storage から GKE 上の vLLM にモデルの重みを読み込むアーキテクチャ。 — vLLM と Cloud Storage を使用した Run:ai Model Streamer アーキテクチャ。

このアーキテクチャには、次のコンポーネントとワークフローが含まれています。

Cloud Storage バケット: AI モデルの重みを safetensors 形式で保存します。
GPU を使用する GKE Pod: vLLM 推論サーバーを実行します。
vLLM 推論サーバー: --load-format=runai_streamer フラグで構成され、モデルストリーマー機能が有効になっています。
Run:ai Model Streamer: vLLM の起動時に、モデルストリーマーは Cloud Storage バケット内の指定された gs:// パスからモデルの重みを読み取ります。ファイルをディスクにダウンロードする代わりに、テンソルデータが GKE Pod の GPU メモリに直接ストリーミングされ、vLLM で推論にすぐに使用できるようになります。
Cloud Storage Anywhere Cache（省略可）: 有効にすると、Anywhere Cache はバケットデータを GKE ノードと同じゾーンにキャッシュに保存し、ストリーマーのデータアクセスをさらに高速化します。

利点

コールドスタート時間の短縮: モデルストリーマーにより、モデルの起動にかかる時間が大幅に短縮されます。従来のメソッドと比較して、モデルの重みを最大 6 倍高速に読み込みます。詳細については、Run:ai モデルストリーマーのベンチマークをご覧ください。
GPU 使用率の向上: モデルの読み込み遅延を最小限に抑えることで、GPU は実際の推論タスクに多くの時間を費やし、全体的な効率と処理能力を高めることができます。
効率化されたワークフロー: このドキュメントで説明するソリューションは GKE と統合されており、vLLM や SGLang などの推論サーバーが Cloud Storage バケット内のモデルに直接アクセスできます。

始める前に

次の前提条件を満たしていることを確認してください。

プロジェクトを選択または作成して API を有効にする

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the Kubernetes Engine, Cloud Storage, Compute Engine, IAM APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the Kubernetes Engine, Cloud Storage, Compute Engine, IAM APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

Cloud Shell を設定する

このドキュメントでは、Google Cloud CLI と kubectl コマンドを使用して、このソリューションに必要なリソースを作成して管理します。 Google Cloud コンソールの上部にある [Cloud Shell をアクティブにする] をクリックすると、これらのコマンドを Cloud Shell で実行できます。

In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

ローカルシェル環境に gcloud CLI をインストールして初期化し、コマンドを実行することもできます。ローカルシェルターミナルを使用する場合は、gcloud auth login コマンドを実行して Google Cloudで認証します。

IAM ロールを付与する

Google Cloud アカウントに、プロジェクトに対する次の IAM ロールがあることを確認します。これにより、GKE クラスタを作成して Cloud Storage を管理できます。

roles/container.admin
roles/storage.admin

これらのロールを付与するには、次のコマンドを実行します。

gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="user:$(gcloud config get-value account)" \
    --role="roles/container.admin"
gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="user:$(gcloud config get-value account)" \
    --role="roles/storage.admin"

PROJECT_ID は、実際のプロジェクト ID に置き換えます。

環境を準備する

このセクションでは、GKE クラスタを設定し、Cloud Storage のモデルにアクセスするための権限を構成する手順について説明します。

GPU を使用して GKE クラスタを作成する

Run:ai Model Streamer は、GKE Autopilot クラスタと Standard クラスタの両方で使用できます。ニーズに最適なクラスタモードを選択してください。

プロジェクト名とクラスタ名の変数を設定します。
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
```
次のように置き換えます。
- PROJECT_ID: 実際の Google Cloud プロジェクト ID。プロジェクト ID は、gcloud config get-value project コマンドを実行して確認できます。
- CLUSTER_NAME: クラスタの名前。例: run-ai-test
Autopilot クラスタまたは Standard クラスタを作成します。
Autopilot
次の手順で GKE Autopilot クラスタを作成します。
1. クラスタのリージョンを設定します。
  export REGION=REGION
  REGION は、クラスタを作成するリージョンに置き換えます。最適なパフォーマンスを得るには、Cloud Storage バケットと同じリージョンを使用します。
2. クラスタを作成します。
  gcloud container clusters create-auto $CLUSTER_NAME \ --project=$PROJECT_ID \ --location=$REGION
Autopilot クラスタは、ワークロードの要件に基づいてノードを自動的にプロビジョニングします。後の手順で vLLM サーバーをデプロイすると、必要に応じて Autopilot が GPU ノードをプロビジョニングします。詳細については、ノードプールの自動作成についてをご覧ください。
標準
GKE Standard クラスタを作成する手順は次のとおりです。
1. クラスタのゾーンを設定します。
  export ZONE=ZONE
  ZONE は、クラスタを作成するゾーンに置き換えます。パフォーマンスを最適化できるよう、Cloud Storage バケットと同じリージョンのゾーンを使用します。
2. クラスタを作成します。
  gcloud container clusters create $CLUSTER_NAME \ --project=$PROJECT_ID \ --zone=$ZONE \ --workload-pool=$PROJECT_ID.svc.id.goog \ --num-nodes=1
3. 1 つの G2 マシン（NVIDIA L4 GPU）を含むノードプールを作成します。
  gcloud container node-pools create g2-gpu-pool \ --cluster=$CLUSTER_NAME \ --zone=$ZONE \ --machine-type=g2-standard-16 \ --num-nodes=1 \ --accelerator=type=nvidia-l4

Workload Identity Federation for GKE を構成する

GKE ワークロードが Cloud Storage バケット内のモデルに安全にアクセスできるように、Workload Identity Federation for GKE を構成します。

Kubernetes サービスアカウントと Namespace の変数を設定します。
```
export KSA_NAME=KSA_NAME
export NAMESPACE=NAMESPACE
```
次のように置き換えます。
- NAMESPACE: ワークロードを実行する Namespace。このドキュメントで作成するすべてのリソースが同じ Namespace になるようにしてください。
- KSA_NAME: Pod が Google Cloud API の認証に使用できる Kubernetes サービスアカウントの名前。
Kubernetes Namespace を作成します。
```
kubectl create namespace $NAMESPACE
```
Kubernetes サービスアカウント（KSA）を作成します。
```
kubectl create serviceaccount $KSA_NAME \
    --namespace=$NAMESPACE
```
注意: 最小権限の原則に準拠するため、同じ Namespace 内の他の Pod でこの Kubernetes サービスアカウントを再利用しないでください。このサービスアカウントを使用するワークロードは、指定された Cloud Storage バケットに意図せずアクセスできてしまいます。
KSA に必要な権限を付与します。
1. 環境変数を設定します。
```
export BUCKET_NAME=BUCKET_NAME
export PROJECT_ID=PROJECT_ID
export PROJECT_NUMBER=$(gcloud projects describe $PROJECT_ID \
    --format 'get(projectNumber)')
```
  次のように置き換えます。
  - BUCKET_NAME: safetensors ファイルを含む Cloud Storage バケットの名前。
  - PROJECT_ID: 実際の Google Cloud プロジェクト ID。
  PROJECT_NUMBER、PROJECT_ID、NAMESPACE、KSA_NAME は、次の手順でプロジェクトの Workload Identity Federation for GKE のプリンシパル ID を作成するために使用されます。
2. Cloud Storage バケット内のオブジェクトを確認できるよう、roles/storage.bucketViewer ロールを KSA に付与します。
```
gcloud storage buckets add-iam-policy-binding gs://$BUCKET_NAME \
    --member="principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/$NAMESPACE/sa/$KSA_NAME" \
    --role="roles/storage.bucketViewer"
```
3. Cloud Storage バケット内のオブジェクトの読み取り、書き込み、削除を行うために、KSA に roles/storage.objectUser ロールを付与します。
```
gcloud storage buckets add-iam-policy-binding gs://$BUCKET_NAME \
    --member="principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/$NAMESPACE/sa/$KSA_NAME" \
    --role="roles/storage.objectUser"
```

GPU を使用する GKE クラスタを設定し、Workload Identity Federation for GKE を構成して、Cloud Storage の AI モデルにアクセスするために必要な権限を Kubernetes のサービスアカウントに付与しました。クラスタと権限を設定したら、vLLM 推論サーバーをデプロイする準備の完了です。このサーバーは、このサービスアカウントを使用して Run:ai Model Streamer でモデルの重みをストリーミングします。

Run:ai Model Streamer を使用して vLLM をデプロイする

vLLM OpenAI 互換サーバーを実行し、Run:ai Model Streamer を使用するように --load-format=runai_streamer フラグで構成された Pod をデプロイします。vLLM のバージョンは 0.11.1 以降である必要があります。

次のマニフェストの例は、単一の NVIDIA L4 GPU を使用して、gemma-2-9b-it などの小規模モデルでモデルストリーマーが有効になっている vLLM 構成になっています。

複数の GPU を必要とする大規模なモデルを使用している場合は、--tensor-parallel-size の値を必要な GPU の数に増やします。
--model-loader-extra-config='{"distributed":true}' フラグは、モデルの重みの分散読み込みを有効にします。これは、オブジェクトストレージからのモデル読み込みパフォーマンスを向上させるためにおすすめの設定です。

詳細については、テンソル並列処理と調整可能なパラメータをご覧ください。

次のマニフェストを vllm-deployment.yaml として保存します。このマニフェストは、Autopilot クラスタと Standard クラスタの両方で柔軟に対応できるように設計されています。

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: vllm-streamer-deployment
      namespace: NAMESPACE
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: vllm-streamer
      template:
        metadata:
          labels:
            app: vllm-streamer
        spec:
          serviceAccountName: KSA_NAME
          containers:
            - name: vllm-container
              image: vllm/vllm-openai:v0.11.1
              command:
                - python3
                - -m
                - vllm.entrypoints.openai.api_server
              args:
                - --model=gs://BUCKET_NAME/PATH_TO_MODEL
                - --load-format=runai_streamer
                - --model-loader-extra-config={"distributed":true}
                - --host=0.0.0.0
                - --port=8000
                - --disable-log-requests
                - --tensor-parallel-size=1
              ports:
                - containerPort: 8000
                  name: api
              # startupProbe allows for longer startup times for large models
              startupProbe:
                httpGet:
                  path: /health
                  port: 8000
                failureThreshold: 60  # 60 * 10s = 10 minutes timeout
                periodSeconds: 10
                initialDelaySeconds: 30
              readinessProbe:
                httpGet:
                  path: /health
                  port: 8000
                failureThreshold: 3
                periodSeconds: 10
              resources:
                limits:
                  nvidia.com/gpu: "1"
                requests:
                  nvidia.com/gpu: "1"
              volumeMounts:
              - mountPath: /dev/shm
                name: dshm
          nodeSelector:
            cloud.google.com/gke-accelerator: nvidia-l4
          volumes:
          - emptyDir:
              medium: Memory
            name: dshm

次のように置き換えます。

NAMESPACE: Kubernetes の Namespace。
KSA_NAME: Kubernetes のサービスアカウント名。
BUCKET_NAME: Cloud Storage バケット名。
PATH_TO_MODEL: バケット内のモデルディレクトリのパス（例: models/my-llama）。

マニフェストを適用して Deployment を作成します。
```
kubectl create -f vllm-deployment.yaml
```

GKE Inference Quickstart ツールを使用すると、より多くの vLLM マニフェストを生成できます。

デプロイを確認する

Deployment のステータスを確認します。
```
kubectl get deployments -n NAMESPACE
```
Pod 名を取得します。
```
kubectl get pods -n NAMESPACE | grep vllm-streamer
```
vllm-streamer-deployment で始まる Pod の名前をメモします。
モデルストリーマーがモデルと重みをダウンロードしたかどうかは、Pod ログで確認できます。
```
kubectl logs -f POD_NAME -n NAMESPACE
```
POD_NAME は、前の手順の Pod の名前に置き換えます。ストリーミングが成功した場合のログは次のようになります。
```
[RunAI Streamer] Overall time to stream 15.0 GiB of all files: 13.4s, 1.1 GiB/s
```

省略可: Anywhere Cache を使用してパフォーマンスを向上させる

Cloud Storage Anywhere Cache を使用すると、GKE ノードに近い場所にデータをキャッシュに保存することで、モデルの読み込みをさらに高速化できます。キャッシュ保存は、同じゾーン内の複数のノードをスケールアウトする場合に特に有効です。

特定の Google Cloud ゾーンの特定の Cloud Storage バケットに対して Anywhere Cache を有効にします。パフォーマンスを向上させるには、キャッシュのゾーンを GKE 推論 Pod が実行されているゾーンと一致させる必要があります。Pod が予測可能なゾーンで実行されるかどうかによって、アプローチが異なります。

GKE Standard ゾーンクラスタで Pod が実行されるゾーンがわかっている場合は、その特定のゾーンに対して Anywhere Cache を有効にします。
複数のゾーンに Pod をスケジュールできるリージョン GKE クラスタ（Autopilot と Standard の両方）では、次のオプションがあります。
- すべてのゾーンでキャッシュ保存を有効にする: クラスタのリージョン内のすべてのゾーンで Anywhere Cache を有効にします。これにより、GKE が Pod をスケジュールする場所に関係なく、キャッシュが使用可能になります。キャッシュ保存が有効になっているゾーンごとに費用が発生します。詳細については、Anywhere Cache の料金をご覧ください。
- Pod を特定のゾーンに配置する: ワークロードマニフェストで nodeSelector ルールまたは nodeAffinity ルールを使用して、Pod を単一ゾーンに制限します。これにより、そのゾーンでのみ Anywhere Cache を有効にできます。ワークロードが単一ゾーンに制限されることを許容できる場合は、この方法がより費用対効果が高くなります。

GKE クラスタが存在するゾーンで Anywhere Cache を有効にするには、次のコマンドを実行します。

# Enable the cache
gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE

# Check the status of the cache
gcloud storage buckets anywhere-caches describe $BUCKET_NAME/$ZONE

クリーンアップ

このドキュメントで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

個々のリソースを削除する手順は次のとおりです。

GKE クラスタを削除します。この操作を行うと、すべてのノードとワークロードが削除されます。
```
gcloud container clusters delete CLUSTER_NAME --location=ZONE_OR_REGION
```
次のように置き換えます。
- CLUSTER_NAME: クラスタの名前。
- ZONE_OR_REGION: クラスタのゾーンまたはリージョン。
Anywhere Cache を有効にしている場合は無効にして、課金が継続的に発生するのを回避します。詳細については、キャッシュを無効にするをご覧ください。

次のステップ

Hugging Face モデルを Cloud Storage に読み込む（試験運用版）方法を確認する。
GPU を使用して GKE で LLM をサービングする方法を確認する。
GKE で利用可能な GPU の詳細を確認する。
GKE のストレージの概要を確認する。
Workload Identity 連携の構成について確認する。

Run:ai Model Streamer と vLLM を使用して GKE でのモデルの読み込みを高速化する

概要

Run:ai Model Streamer

safetensors

vLLM

Run:ai Model Streamer がモデルの読み込みを高速化する仕組み

アーキテクチャの概要

利点

始める前に

プロジェクトを選択または作成して API を有効にする

Cloud Shell を設定する

IAM ロールを付与する

環境を準備する

GPU を使用して GKE クラスタを作成する

Autopilot

標準

Workload Identity Federation for GKE を構成する

Run:ai Model Streamer を使用して vLLM をデプロイする

デプロイを確認する

省略可: Anywhere Cache を使用してパフォーマンスを向上させる

クリーンアップ

次のステップ