このガイドでは、Google Kubernetes Engine(GKE)で Tensor Processing Unit(TPU)を使用し、JetStream と PyTorch で大規模言語モデル(LLM)をサービングする方法について説明します。このガイドでは、モデルの重みを Cloud Storage にダウンロードし、JetStream を実行するコンテナを使用して GKE Autopilot または Standard クラスタにデプロイします。
モデルを JetStream にデプロイするときに Kubernetes の機能によって実現されるスケーラビリティ、復元力、費用対効果が必要な場合は、このガイドが適しています。
このガイドは、PyTorch を使用している生成 AI をご利用のお客様、GKE の新規または既存のユーザー、ML エンジニア、MLOps(DevOps)エンジニア、LLM のサービングに Kubernetes コンテナのオーケストレーション機能を使用することに関心をお持ちのプラットフォーム管理者を対象としています。
背景
GKE で TPU を使用して JetStream で LLM をサービングすることで、マネージド Kubernetes のメリット(費用効率、スケーラビリティ、高可用性など)をすべて活用した、プロダクション レディな堅牢なサービング ソリューションを構築できます。このセクションでは、このチュートリアルで使用されている重要なテクノロジーについて説明します。
TPU について
TPU は、Google が独自に開発した特定用途向け集積回路(ASIC)であり、TensorFlow、PyTorch、JAX などのフレームワークを使用して構築された ML モデルと AI モデルを高速化するために使用されます。
GKE で TPU を使用する前に、次の学習プログラムを完了することをおすすめします。
- Cloud TPU システム アーキテクチャで、現在の TPU バージョンの可用性について学習する。
- GKE の TPU についてを確認する。
このチュートリアルでは、さまざまな LLM モデルのサービングについて説明します。GKE は、低レイテンシでプロンプトをサービングするモデルの要件に基づいて構成された TPU トポロジを使用して、単一ホストの TPUv5e ノードにモデルをデプロイします。
JetStream について
JetStream は、Google が開発したオープンソースの推論サービング フレームワークです。JetStream を使用すると、TPU と GPU で高性能、高スループット、メモリ最適化された推論が可能になります。JetStream では、連続バッチ処理、KV キャッシュの最適化、量子化手法などの高度なパフォーマンス最適化により、LLM を簡単にデプロイできます。JetStream では、PyTorch / XLA と JAX TPU のサービングにより、最適なパフォーマンスを実現できます。
連続的なバッチ処理
連続的バッチ処理は、受信した推論リクエストを動的にバッチにグループ化し、レイテンシを短縮してスループットを向上させる手法です。
KV キャッシュの量子化
KV キャッシュの量子化では、アテンション メカニズムで使用される Key-Value キャッシュを圧縮して、メモリ要件を削減します。
Int8 重み量子化
Int8 重み量子化では、モデル重みの精度を 32 ビットの浮動小数点数から 8 ビットの整数にすることで、計算速度を向上させ、メモリ使用量を削減しています。
これらの最適化の詳細については、JetStream PyTorch と JetStream MaxText のプロジェクト リポジトリをご覧ください。
PyTorch について
PyTorch は、Meta によって開発されたオープンソースの ML フレームワークで、現在は Linux Foundation 傘下にあります。PyTorch は、テンソル計算やディープ ニューラル ネットワークなどの高度な機能を提供します。
目標
- モデルの特性に基づいて推奨される TPU トポロジを持つ GKE Autopilot または Standard クラスタを準備します。
- GKE に JetStream コンポーネントをデプロイします。
- モデルを取得して公開します。
- 公開されたモデルをサービングして操作します。
アーキテクチャ
このセクションでは、このチュートリアルで使用する GKE アーキテクチャについて説明します。このアーキテクチャには、TPU をプロビジョニングし、モデルをデプロイしてサービングするための JetStream コンポーネントをホストする GKE Autopilot または Standard クラスタが含まれています。
次の図は、このアーキテクチャのコンポーネントを示しています。
このアーキテクチャには次のコンポーネントが含まれています。
- GKE Autopilot または Standard リージョン クラスタ。
- JetStream デプロイをホストする 2 つの単一ホスト TPU スライス ノードプール。
- Service コンポーネントは、インバウンド トラフィックをすべての
JetStream HTTP
レプリカに分散します。 JetStream HTTP
は、JetStream の必須フォーマットのラッパーとしてリクエストを受け取り、JetStream の GRPC クライアントに送信する HTTP サーバーです。JetStream-PyTorch
は、継続的バッチ処理で推論を行う JetStream サーバーです。
始める前に
- 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.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required API.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required API.
-
Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin
Check for the roles
-
In the Google Cloud console, go to the IAM page.
Go to IAM - Select the project.
-
In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.
- For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.
Grant the roles
-
In the Google Cloud console, go to the IAM page.
[IAM] に移動 - プロジェクトを選択します。
- [ アクセスを許可] をクリックします。
-
[新しいプリンシパル] フィールドに、ユーザー ID を入力します。 これは通常、Google アカウントのメールアドレスです。
- [ロールを選択] リストでロールを選択します。
- 追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
- [保存] をクリックします。
-
- 8 個の TPU v5e PodSlice Lite チップに十分な割り当てがあることを確認します。このチュートリアルでは、オンデマンド インスタンスを使用します。
- Hugging Face トークンを作成します(まだ作成していない場合)。
モデルへのアクセス権を取得する
GKE へのデプロイ用に Hugging Face のさまざまなモデルにアクセスする
Gemma 7B-it
GKE にデプロイするために Gemma モデルへのアクセス権を取得するには、まずライセンス同意契約に署名する必要があります。
- Hugging Face の Gemma モデルの同意ページにアクセスします。
- Hugging Face にログインします(まだログインしていない場合)。
- モデルの利用規約を確認して同意します。
Llama 3 8B
GKE にデプロイするために Llama 3 モデルへのアクセス権を取得するには、まず、ライセンス同意契約に署名する必要があります。
- Hugging Face の Llama 3 モデルの同意ページにアクセスします。
- Hugging Face にログインします(まだログインしていない場合)。
- モデルの利用規約を確認して同意します。
環境を準備する
このチュートリアルでは、Cloud Shell を使用して Google Cloud でホストされるリソースを管理します。Cloud Shell には、このチュートリアルに必要な kubectl
や gcloud CLI などのソフトウェアがプリインストールされています。
Cloud Shell を使用して環境を設定するには、次の操作を行います。
Google Cloud コンソールで、Google Cloud コンソールの [Cloud Shell をアクティブにする] をクリックして、Cloud Shell セッションを起動します。これにより、Google Cloud コンソールの下部ペインでセッションが起動されます。
デフォルトの環境変数を設定します。
gcloud config set project PROJECT_ID export PROJECT_ID=$(gcloud config get project) export CLUSTER_NAME=CLUSTER_NAME export BUCKET_NAME=BUCKET_NAME export REGION=REGION export LOCATION=LOCATION export CLUSTER_VERSION=CLUSTER_VERSION
次の値を置き換えます。
PROJECT_ID
: Google Cloud のプロジェクト ID。CLUSTER_NAME
: GKE クラスタの名前。BUCKET_NAME
: Cloud Storage バケットの名前。gs://
プレフィックスを指定する必要はありません。REGION
: GKE クラスタ、Cloud Storage バケット、TPU ノードが配置されているリージョン。これは、TPU v5e マシンタイプを使用できるゾーン(us-west1
、us-west4
、us-central1
、us-east1
、us-east5
、europe-west4
など)が含まれているリージョンです。Autopilot クラスタの場合は、選択したリージョンに十分な TPU v5e ゾーンリソースがあることを確認してください。- (Standard クラスタのみ)
LOCATION
: TPU リソースが使用可能なゾーン(us-west4-a
など)。Autopilot クラスタの場合は、リージョンのみを指定します。ゾーンを指定する必要はありません。 CLUSTER_VERSION
: GKE バージョン。使用するマシンタイプをサポートしている必要があります。デフォルトの GKE バージョンは、ターゲット TPU で利用できない場合があります。TPU マシンタイプで使用できる最小 GKE バージョンのリストについては、GKE での TPU の可用性をご覧ください。
Google Cloud リソースを作成して構成する
次の手順に沿って、必要なリソースを作成します。
GKE クラスタを作成する
GKE Autopilot クラスタまたは GKE Standard クラスタの GPU で Gemma をサービングできます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な GKE の運用モードを選択するには、GKE の運用モードを選択するをご覧ください。
Autopilot
Autopilot GKE クラスタを作成します。
gcloud container clusters create-auto CLUSTER_NAME \
--project=PROJECT_ID \
--region=REGION \
--cluster-version=CLUSTER_VERSION
Standard
Workload Identity Federation for GKE を使用するリージョン GKE Standard クラスタを作成します。
gcloud container clusters create CLUSTER_NAME \ --enable-ip-alias \ --machine-type=e2-standard-4 \ --num-nodes=2 \ --cluster-version=CLUSTER_VERSION \ --workload-pool=PROJECT_ID.svc.id.goog \ --location=REGION
クラスタの作成には数分かかることもあります。
2x4
トポロジと 2 つのノードを持つ TPU v5e ノードプールを作成します。gcloud container node-pools create tpu-nodepool \ --cluster=CLUSTER_NAME \ --machine-type=ct5lp-hightpu-8t \ --project=PROJECT_ID \ --num-nodes=2 \ --region=REGION \ --node-locations=LOCATION
Cloud Storage バケットを作成する
変換されたチェックポイントを保存する Cloud Storage バケットを作成します。
gcloud storage buckets create gs://BUCKET_NAME --location=REGION
Cloud Shell で Hugging Face CLI トークンを生成する
Hugging Face トークンをまだ生成していない場合は、新しいトークンを生成します。
- [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
- [New Token] をクリックします。
- 任意の名前と、少なくとも
Read
ロールを指定します。 - [Generate a token] をクリックします。
- モデルの Hugging Face リポジトリに対する読み取りアクセス権を取得するには、アクセス トークンの権限を編集します。
- トークンをクリップボードにコピーします。
Hugging Face の認証情報用の Kubernetes Secret を作成する
Cloud Shell で、次の操作を行います。
クラスタと通信を行うように
kubectl
を構成します。gcloud container clusters get-credentials CLUSTER_NAME --location=REGION
Hugging Face の認証情報を保存する Secret を作成します。
kubectl create secret generic huggingface-secret \ --from-literal=HUGGINGFACE_TOKEN=HUGGINGFACE_TOKEN
HUGGINGFACE_TOKEN
は、Hugging Face トークンに置き換えます。
Workload Identity Federation for GKE を使用してワークロード アクセスを構成する
アプリケーションに Kubernetes ServiceAccount を割り当て、IAM サービス アカウントとして機能するようにその Kubernetes ServiceAccount を構成します。
アプリケーションの IAM サービス アカウントを作成します。
gcloud iam service-accounts create wi-jetstream
Cloud Storage を管理する IAM サービス アカウントの IAM ポリシー バインディングを追加します。
gcloud projects add-iam-policy-binding PROJECT_ID \ --member "serviceAccount:wi-jetstream@PROJECT_ID.iam.gserviceaccount.com" \ --role roles/storage.objectUser gcloud projects add-iam-policy-binding PROJECT_ID \ --member "serviceAccount:wi-jetstream@PROJECT_ID.iam.gserviceaccount.com" \ --role roles/storage.insightsCollectorService
2 つのサービス アカウントの間に IAM ポリシー バインディングを追加して、Kubernetes ServiceAccount が IAM サービス アカウントの権限を借用できるようにします。このバインドにより、Kubernetes ServiceAccount が IAM サービス アカウントとして機能します。
gcloud iam service-accounts add-iam-policy-binding wi-jetstream@PROJECT_ID.iam.gserviceaccount.com \ --role roles/iam.workloadIdentityUser \ --member "serviceAccount:PROJECT_ID.svc.id.goog[default/default]"
Kubernetes サービス アカウントに IAM サービス アカウントのメールアドレスでアノテーションを付けます。
kubectl annotate serviceaccount default \ iam.gke.io/gcp-service-account=wi-jetstream@PROJECT_ID.iam.gserviceaccount.com
JetStream をデプロイする
JetStream コンテナをデプロイしてモデルをサービングします。
次のマニフェストを jetstream-pytorch-deployment.yaml
として保存します。
Gemma 7B-it
Llama 3 8B
このマニフェストでは、次の主要なプロパティが設定されています。
model_id
: Hugging Face のモデル名(google/gemma-7b-it
、meta-llama/Meta-Llama-3-8B
)(サポートされているモデルをご覧ください)。override_batch_size
: デバイスごとのデコード バッチサイズ。1 つの TPU チップは 1 つのデバイスに相当します。この値のデフォルトは32
です。working_dir
: モデル チェックポイントが保存されている、または保存される作業ディレクトリ。Cloud Storage バケットがノードにマウントされているため、ワークロードはgs://BUCKET_NAME/pytorch/org/repo
のチェックポイントにアクセスできます。JetStream-PyTorch サーバーは、初回使用時に Hugging Face からモデルの重みを作業ディレクトリにダウンロードします。70 億または 80 億のモデルの読み込みには数分かかることがあります。その後のデプロイでは、JetStream-PyTorch サーバーは作業ディレクトリからモデル重みを読み込み、読み込み時間が短縮されます。既存のチェックポイントの場合は、作業ディレクトリで同じパス規則を使用していることを確認します。enable_model_warmup
: この設定により、モデルサーバーの起動後にモデルのウォームアップが有効になります。この値のデフォルトはFalse
です。
必要に応じて、次のプロパティを設定できます。
max_input_length
: 最大入力シーケンス長。この値のデフォルトは1024
です。max_output_length
: 最大出力デコード長。デフォルト値は1024
です。quantize_weights
: チェックポイントが量子化されているかどうか。この値のデフォルトは0
です。int8
量子化を有効にするには、この値を 1 に設定します。internal_jax_compilation_cache
: JAX コンパイル キャッシュのディレクトリ。この値のデフォルトは~/jax_cache
です。リモート キャッシュの場合はgs://BUCKET_NAME/jax_cache
に設定します。
マニフェストでは、モデルが読み込まれてウォームアップが完了した後にモデルサーバーに Ready
というラベルが付けられるように、起動プローブが構成されています。モデルサーバーの健全性を確保するために、liveness プローブと readiness プローブが構成されています。
BUCKET_NAME
は、前に作成した GSBucket に置き換えます。sed -i "s|BUCKET_NAME|BUCKET_NAME|g" jetstream-pytorch-deployment.yaml
次のようにマニフェストを適用します。
kubectl apply -f jetstream-pytorch-deployment.yaml
Deployment を確認します。
kubectl get deployment
出力は次のようになります。
NAME READY UP-TO-DATE AVAILABLE AGE jetstream-pytorch-server 0/2 2 0 ##s
Autopilot クラスタの場合、必要な TPU リソースのプロビジョニングに数分かかることがあります。
JetStream-PyTorch サーバーログを表示して、モデルの重みが読み込まれ、モデルのウォームアップが完了したことを確認します。このオペレーションが完了するまでに数分かかることがあります。
kubectl logs deploy/jetstream-pytorch-server -f -c jetstream-pytorch-server
出力は次のようになります。
Started jetstream_server.... 2024-04-12 04:33:37,128 - root - INFO - ---------Generate params 0 loaded.---------
Deployment の準備ができていることを確認します。
kubectl get deployment
出力は次のようになります。
NAME READY UP-TO-DATE AVAILABLE AGE jetstream-pytorch-server 2/2 2 2 ##s
healthcheck
エンドポイントの登録には数分かかることがあります。
モデルをサービングする
このセクションでは、モデルを操作します。
ポート転送をセットアップする
JetStream Deployment には、前の手順で作成した ClusterIP Service を介してアクセスできます。ClusterIP Service にはクラスタ内からのみアクセスできます。したがって、クラスタの外部から Service にアクセスするには、次の手順を完了します。
ポート転送セッションを確立するには、次のコマンドを実行します。
kubectl port-forward svc/jetstream-svc 8000:8000
curl を使用してモデルを操作する
新しいターミナルを開いて次のコマンドを実行し、JetStream HTTP サーバーにアクセスできることを確認します。
curl --request POST \ --header "Content-type: application/json" \ -s \ localhost:8000/generate \ --data \ '{ "prompt": "What are the top 5 programming languages", "max_tokens": 200 }'
モデルのウォームアップにより、最初のリクエストが完了するまでに数秒かかることがあります。出力は次のようになります。
{ "response": " for data science in 2023?\n\n**1. Python:**\n- Widely used for data science due to its readability, extensive libraries (pandas, scikit-learn), and integration with other tools.\n- High demand for Python programmers in data science roles.\n\n**2. R:**\n- Popular choice for data analysis and visualization, particularly in academia and research.\n- Extensive libraries for statistical modeling and data wrangling.\n\n**3. Java:**\n- Enterprise-grade platform for data science, with strong performance and scalability.\n- Widely used in data mining and big data analytics.\n\n**4. SQL:**\n- Essential for data querying and manipulation, especially in relational databases.\n- Used for data analysis and visualization in various industries.\n\n**5. Scala:**\n- Scalable and efficient for big data processing and machine learning models.\n- Popular in data science for its parallelism and integration with Spark and Spark MLlib." }
次の作業が完了しました。
- TPU を使用して、GKE に JetStream-PyTorch モデル サーバーをデプロイしました。
gs://BUCKET_NAME/pytorch/org/repo
にチェックポイントを作成しました。- モデルをサービングして操作しました。
問題のトラブルシューティング
Empty reply from server
というメッセージが表示された場合は、コンテナがモデルデータのダウンロードを完了していない可能性があります。モデルがサービス提供の準備ができていることを示すConnected
というメッセージがないか、再度 Pod のログを確認します。Connection refused
が表示された場合は、ポート転送が有効であることを確認します。
クリーンアップ
このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。
デプロイされたリソースを削除する
このガイドで作成したリソースに対して Google Cloud アカウントで課金されないようにするには、次のコマンドを実行し、プロンプトに従います。
gcloud container clusters delete CLUSTER_NAME --region=REGION
gcloud iam service-accounts delete wi-jetstream@PROJECT_ID.iam.gserviceaccount.com
gcloud storage rm --recursive gs://BUCKET_NAME
次のステップ
- GKE で Gemma モデルを実行する方法と、GKE プラットフォームのオーケストレーション機能を使用して最適化された AI / ML ワークロードを実行する方法を確認する。
- GKE の TPU の詳細を確認する。
- JetStream の GitHub リポジトリを確認する。
- Vertex AI Model Garden を確認する。