Saxml を備えた Gemini Enterprise エージェント プラットフォームで TPU を使用して Gemma オープンモデルを提供する

このガイドでは、Saxml と Gemini Enterprise Agent Platform で Tensor Processing Unit(TPU)を使用して Gemma オープンモデルの大規模言語モデル(LLM)をサービングする方法について説明します。このガイドでは、2B および 7B パラメータ指示でチューニングした Gemma モデルを Cloud Storage にダウンロードし、TPU で Saxml を実行する Gemini Enterprise Agent Platform にデプロイします。

背景

Saxml と Gemini Enterprise Agent Platform で TPU を使用して Gemma をサービングします。低レベルのインフラストラクチャを処理し、費用対効果の高い方法で LLM をサービングできるマネージド AI ソリューションを利用できます。このセクションでは、このチュートリアルで使用されている重要なテクノロジーについて説明します。

Gemma

Gemma は、オープン ライセンスでリリースされ一般公開されている、軽量の生成 AI モデルのセットです。これらの AI モデルは、アプリケーション、ハードウェア、モバイル デバイス、ホスト型サービスで実行できます。Gemma モデルはテキスト生成に使用できますが、特殊なタスク用にチューニングすることもできます。

詳しくは、Gemma のドキュメントをご覧ください。

Saxml

Saxml は、推論に Paxml、JAX、PyTorch モデルを提供する試験運用版のシステムです。 このチュートリアルでは、Saxml のコスト効率の高い TPU で Gemma を提供する方法について説明します。GPU の設定も同様です。Saxml は、このチュートリアルで使用する Gemini Enterprise Agent Platform のコンテナをビルドするスクリプトを提供します。

TPU

TPU は、Google が独自に開発した特定用途向け集積回路(ASIC)であり、TensorFlow、PyTorch、JAX などのデータ処理フレームワークを高速化するために使用されます。

このチュートリアルでは、Gemma 2B モデルと Gemma 7B モデルを使用します。Gemini Enterprise Agent Platform は、次のシングルホスト TPU v5e ノードプールでこれらのモデルをホストします。

  • Gemma 2B: 1 つの TPU チップを表す 1x1 トポロジを持つ TPU v5e ノードプールでホストされます。ノードのマシンタイプは ct5lp-hightpu-1t です。
  • Gemma 7B: 4 つの TPU チップを表す 2x2 トポロジを持つ TPU v5e ノードプールでホストされます。ノードのマシンタイプは ct5lp-hightpu-4t です。

始める前に

  1. In the Cloud de Confiance console, on the project selector page, select or create a Cloud de Confiance 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 role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  2. Verify that billing is enabled for your Cloud de Confiance project.

  3. Enable the Agent Platform API, if it is not already enabled.

    Roles required to enable APIs

    To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

    Enable the API

  4. Cloud de Confiance コンソールで Cloud Shell をアクティブにします。

    Cloud Shell をアクティブにする

    Cloud de Confiance コンソールの下部で Cloud Shell セッションが開始し、コマンドライン プロンプトが表示されます。Cloud Shell はシェル環境です。Google Cloud CLI がすでにインストールされており、現在のプロジェクトの値もすでに設定されています。セッションが初期化されるまで数秒かかることがあります。

このチュートリアルでは、 Cloud de Confiance by S3NSの操作に Cloud Shell を使用していることを前提としています。Cloud Shell の代わりに別のシェルを使用する場合は、次の追加の構成を行います。

  1. Google Cloud CLI をインストールします。

  2. フェデレーション ID(連携 ID)を使用するように gcloud CLI を構成します。

    詳細については、 フェデレーション ID を使用して gcloud CLI にログインするをご覧ください。

  3. gcloud CLI を初期化するには、次のコマンドを実行します。

    gcloud init
  4. Gemini Enterprise Agent Platform の TPU v5e チップに十分な割り当てがあることを確認します。デフォルトでは、この割り当ては 0 です。1x1 トポロジの場合は 1 にする必要があります。2x2 の場合は 4 にする必要があります。両方のトポロジを実行するには、5 にする必要があります。
  5. Kaggle アカウントを作成します(まだアカウントを保有されていない場合)。

モデルへのアクセス権を取得する

Cloud Shell にモデルの重みをダウンロードするのに十分なリソースがない場合があります。その場合は、そのタスクを実行する Gemini Enterprise Agent Platform Workbench インスタンスを作成できます。

Gemini Enterprise Agent Platform にデプロイするために Gemma モデルへのアクセス権を取得するには、Kaggle プラットフォームにログインし、ライセンス同意契約に署名して、Kaggle API トークンを入手する必要があります。このチュートリアルでは、Kaggle 認証情報に Kubernetes Secret を使用します。

Gemma を使用するには同意契約に署名する必要があります。手順は次のとおりです。

  1. Kaggle.com のモデルの同意ページにアクセスします。
  2. Kaggle にログインしていない場合はログインします。
  3. [アクセス権限をリクエスト] をクリックします。
  4. [同意に使用するアカウントを選択] セクションで、[Kaggle アカウントを使用して確認] を選択して、同意に Kaggle アカウントを使用します。
  5. モデルの利用規約に同意します。

アクセス トークンを生成する

Kaggle からモデルにアクセスするには、Kaggle API トークンが必要です。

トークンをまだ生成していない場合は、次の手順に沿って生成します。

  1. ブラウザで [Kaggle の設定] に移動します。
  2. [API] セクションで [新しいトークンを作成] をクリックします。

    kaggle.json という名前のファイルがダウンロードされます。

アクセス トークンを Cloud Shell にアップロードする

Cloud Shell で、Kaggle API トークンを Cloud de Confiance by S3NSプロジェクトにアップロードできます。

  1. Cloud Shell で、 [その他] > [アップロード] をクリックします。
  2. [ファイル] を選択し、[ファイルを選択] をクリックします。
  3. kaggle.json ファイルを開きます。
  4. [アップロード] をクリックします。

Cloud Storage バケットを作成する

モデルのチェックポイントを保存する Cloud Storage バケットを作成する。

Cloud Shell で次のコマンドを実行します。

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

CHECKPOINTS_BUCKET_NAME は、モデルのチェックポイントを保存する Cloud Storage バケットの名前に置き換えます。

モデルを Cloud Storage バケットにコピーする

Cloud Shell で次のコマンドを実行します。

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

モデルのデプロイ

モデルをアップロードする

Saxml コンテナを使用する Model リソースをアップロードするには、次の gcloud ai models upload コマンドを実行します。

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

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

  • PROJECT_ID: Cloud de Confianceプロジェクトの ID。
  • LOCATION_ID: Agent Platform を使用しているリージョン。TPU は us-west1 でのみ使用できます。
  • DEPLOYED_MODEL_NAME: DeployedModel の名前。DeployedModel の Model の表示名を使用することもできます。

エンドポイントを作成する

モデルを使用してオンライン推論を行う前に、モデルをエンドポイントにデプロイする必要があります。既存のエンドポイントにモデルをデプロイする場合は、この手順を省略できます。次の例では、gcloud ai endpoints create コマンドを使用します。

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

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

  • LOCATION_ID: Agent Platform を使用しているリージョン。
  • ENDPOINT_NAME: エンドポイントの表示名。

Google Cloud CLI ツールがエンドポイントを作成するまでに数秒かかる場合があります。

モデルをエンドポイントにデプロイする

エンドポイントの準備が整ったら、モデルをエンドポイントにデプロイします。

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

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

  • LOCATION_ID: Agent Platform を使用しているリージョン。
  • ENDPOINT_NAME: エンドポイントの表示名。
  • DEPLOYED_MODEL_NAME: DeployedModel の名前。DeployedModel の Model の表示名を使用することもできます。

Gemma 2B はより小さい ct5lp-hightpu-1t マシンにデプロイできます。この場合、モデルをアップロードするときに --platform_topology=1x1 を指定する必要があります。

Google Cloud CLI ツールでエンドポイントにモデルをデプロイする場合、処理に数分かかることがあります。モデルが正常にデプロイされると、このコマンドは次の出力を返します。

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

デプロイされたモデルからオンライン推論を取得する

Gemini Enterprise Agent Platform エンドポイントからモデルを呼び出すには、標準の推論リクエストの JSON オブジェクト を使用して推論リクエストをフォーマットします。

次の例では、gcloud ai endpoints predict コマンドを使用します。

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

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

  • LOCATION_ID: Agent Platform を使用しているリージョン。
  • ENDPOINT_NAME: エンドポイントの表示名。
  • instances.json の形式は {"instances": [{"text_batch": "<your prompt>"},{...}]} です。

クリーンアップ

Gemini Enterprise Agent Platform の料金と Artifact Registry の料金が発生しないように、このチュートリアルで作成した Cloud de Confiance リソースを削除します。

  1. エンドポイントからモデルのデプロイを解除してエンドポイントを削除するには、シェルで次のコマンドを実行します。

    ENDPOINT_ID=$(gcloud ai endpoints list \
       --region=LOCATION \
       --filter=display_name=ENDPOINT_NAME \
       --format="value(name)")
    
    DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
       --region=LOCATION \
       --format="value(deployedModels.id)")
    
    gcloud ai endpoints undeploy-model $ENDPOINT_ID \
      --region=LOCATION \
      --deployed-model-id=$DEPLOYED_MODEL_ID
    
    gcloud ai endpoints delete $ENDPOINT_ID \
       --region=LOCATION \
       --quiet
    

    LOCATION は、前のセクションでモデルを作成したリージョンに置き換えます。

  2. モデルを削除するには、シェルで次のコマンドを実行します。

    MODEL_ID=$(gcloud ai models list \
       --region=LOCATION \
       --filter=display_name=DEPLOYED_MODEL_NAME \
       --format="value(name)")
    
    gcloud ai models delete $MODEL_ID \
       --region=LOCATION \
       --quiet
    

    LOCATION は、前のセクションでモデルを作成したリージョンに置き換えます。

制限事項

  • Gemini Enterprise Agent Platform では、Cloud TPU は us-west1 でのみサポートされています。詳細については、ロケーションをご覧ください。

次のステップ

  • Llama2 や GPT-J など、他の Saxml モデルをデプロイする方法について確認する。