TPU を使用してパイプラインを実行する

このページでは、TPU を使用して Dataflow で Apache Beam パイプラインを実行する方法について説明します。TPU を使用するジョブでは、Dataflow の料金ページで定められた料金が発生します。

Dataflow で TPU を使用する方法については、Dataflow での TPU のサポートをご覧ください。

省略可: アクセラレータを使用するための特定の予約を行う

TPU はオンデマンドで使用できますが、特定のターゲット Google Cloud予約で Dataflow TPU を使用することを強くおすすめします。これにより、利用可能なアクセラレータにアクセスでき、ワーカーの起動時間を短縮できます。TPU 予約を使用するパイプラインには、追加の TPU 割り当ては必要ありません。

予約を行わずに TPU をオンデマンドで使用する場合は、パイプラインを実行する前に TPU 割り当てをプロビジョニングします。

省略可: TPU 割り当てをプロビジョニングする

TPU は、オンデマンド容量または予約を使用して使用できます。TPU をオンデマンドで使用する場合は、事前に TPU 割り当てをプロビジョニングする必要があります。特定のターゲットの予約を使用する場合は、このセクションをスキップできます。

予約なしで TPU をオンデマンドで使用するには、次の手順で TPU 用の Compute Engine API の割り当ての上限と現在の使用量を確認します。

コンソール

  1. Google Cloud コンソールで [割り当て] ページに移動します。

    [割り当て] に移動

  2. [ フィルタ] ボックスで、次の操作を行います。

    1. 次の表を使用して、TPU のバージョンとマシンタイプに基づいて割り当てのプロパティを選択してコピーします。たとえば、マシンタイプが「ct5lp-」で始まるオンデマンド TPU v5e ノードを作成する場合は、「Name: TPU v5 Lite PodSlice chips」と入力します。

      TPU バージョン、次で始まるマシンタイプ オンデマンド インスタンスの割り当てのプロパティと名前
      TPU v5e、
      ct5lp- 
      Name:
      TPU v5 Lite PodSlice chips
      TPU v5p、
      ct5p- 
      Name:
      TPU v5p chips
      TPU v6e、
      ct6e- 
      Dimensions (e.g. location):
      tpu_family:CT6E

    2. [ディメンション（ロケーションなど）] プロパティを選択し、「region:」に続けて、パイプラインを開始するリージョンの名前を入力します。たとえば、ゾーン us-west4-a を使用する場合は、「region:us-west4」と入力します。TPU の割り当てはリージョン単位であるため、同じリージョン内のすべてのゾーンで同じ TPU の割り当てが消費されます。

カスタム コンテナ イメージを構成する

Dataflow パイプラインで TPU を操作するには、パイプラインのランタイム環境で XLA デバイスを操作できるソフトウェアを提供する必要があります。これには、パイプラインのニーズに基づいて TPU ライブラリをインストールし、使用する TPU デバイスに基づいて環境変数を構成する必要があります。

コンテナ イメージをカスタマイズするには、必要な TPU ライブラリを含む既製のベースイメージに Apache Beam をインストールします。または、Apache Beam SDK リリースで公開されたイメージに TPU ソフトウェアをインストールします。

カスタム コンテナ イメージを指定するには、sdk_container_image パイプライン オプションを使用します。詳細については、Dataflow でカスタム コンテナを使用するをご覧ください。

TPU アクセラレータを使用する場合は、コンテナ イメージで次の環境変数を設定する必要があります。

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

使用するアクセラレータに応じて、次の表の変数も設定する必要があります。

type トポロジ Dataflow worker_machine_type が必須になりました 追加の環境変数
tpu-v5-lite-podslice 1×1 ct5lp-hightpu-1t 
TPU_ACCELERATOR_TYPE=v5litepod-1
TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t 
TPU_ACCELERATOR_TYPE=v5litepod-4
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t 
TPU_ACCELERATOR_TYPE=v5litepod-8
TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice 1×1 ct6e-standard-1t 
TPU_ACCELERATOR_TYPE=v6e-1
TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice 2x2 ct6e-standard-4t 
TPU_ACCELERATOR_TYPE=v6e-4
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice 2x4 ct6e-standard-8t 
TPU_ACCELERATOR_TYPE=v6e-8
TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t 
TPU_ACCELERATOR_TYPE=v5p-8
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

カスタム コンテナ イメージのサンプル Dockerfile は、次の例のようになります。

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

TPU を使用してジョブを実行する

TPU を使用した Dataflow ジョブを実行する際の考慮事項は次のとおりです。

  • TPU コンテナは大きくなる可能性があるため、ディスク容量が不足しないように、--disk_size_gb パイプライン オプションを使用して、デフォルトのブートディスク サイズを 50 GB またはコンテナ イメージで必要な適切なサイズに増やします。
  • ワーカー内の並列処理を制限します。

TPU とワーカーの並列処理

デフォルトの構成では、Dataflow Python パイプラインは VM コアごとに 1 つの Apache Beam SDK プロセスを起動します。TPU マシンタイプには多数の vCPU コアがありますが、TPU デバイスで計算を実行できるのは 1 つのプロセスのみです。また、TPU デバイスは、プロセスの存続期間中、プロセスによって予約されることがあります。したがって、Dataflow TPU パイプラインを実行する場合は、ワーカー内の並列処理を制限する必要があります。ワーカーの並列処理を制限するには、次のガイダンスを使用します。

  • ユースケースでモデルに対する推論の実行が必要な場合は、Beam RunInference API を使用します。詳細については、Beam での大規模言語モデルの推論をご覧ください。
  • Beam RunInference API を使用できない場合は、Beam のマルチプロセス共有オブジェクトを使用して、特定のオペレーションを単一のプロセスに制限します。
  • 上記の推奨事項を使用できず、ワーカーごとに 1 つの Python プロセスのみを起動する場合は、--experiments=no_use_multiple_sdk_containers パイプライン オプションを設定します。
  • 100 個を超える vCPU を持つワーカーの場合は、--number_of_worker_harness_threads パイプライン オプションを使用してスレッド数を減らします。次の表を使用して、TPU タイプが 100 個を超える vCPU を使用しているかどうかを確認します。

次の表に、各 TPU 構成のワーカーあたりの合計コンピューティング リソースを示します。

TPU タイプ トポロジ マシンタイプ TPU チップ vCPU RAM（GB）
tpu-v5-lite-podslice 1×1 ct5lp-hightpu-1t 1 24 48
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t 4 112 192
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t 8 224 384
tpu-v6e-slice 1×1 ct6e-standard-1t 1 44 176
tpu-v6e-slice 2x2 ct6e-standard-4t 4 180 720
tpu-v6e-slice 2x4 ct6e-standard-8t 8 360 1440
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t 4 208 448

TPU を使用してパイプラインを実行する

TPU を使用して Dataflow ジョブを実行するには、次のコマンドを使用します。

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

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

  • PIPELINE: パイプラインのソースコード ファイル。
  • PROJECT: Google Cloud プロジェクト名。
  • BUCKET: Cloud Storage バケット。
  • REGION: Dataflow リージョン（例: us-central1）。
  • TPU_TYPE: サポートされている TPU タイプ（tpu-v5-lite-podslice など）。タイプとトポロジの完全なリストについては、サポートされている TPU アクセラレータをご覧ください。
  • TPU_TOPOLOGY: TPU トポロジ（1x1 など）。
  • MACHINE_TYPE: 対応するマシンタイプ（例: ct5lp-hightpu-1t）。
  • DISK_SIZE_GB: 各ワーカー VM のブートディスクのサイズ（例: 100）。
  • IMAGE: Docker イメージの Artifact Registry パス。
  • NUMBER_OF_THREADS: 省略可。ワーカー ハンドラ スレッドの数。

Dataflow ジョブを確認する

ジョブが TPU を持つワーカー VM を使用していることを確認する手順は次のとおりです。

  1. Google Cloud コンソールで、[Dataflow] > [ジョブ] ページに移動します。

    [ジョブ] に移動

  2. ジョブを選択します。

  3. [ジョブの指標] タブをクリックします。

  4. [自動スケーリング] セクションで、現在のワーカー VM が 1 つ以上あることを確認します。

  5. [ジョブ情報] サイドペインで、machine_type が ct で始まることを確認します。たとえば、ct6e-standard-1t のようにします。これは TPU の使用率を示します。

Dataflow ジョブのトラブルシューティング

TPU で Dataflow ジョブを実行する際に問題が発生した場合は、Dataflow TPU ジョブのトラブルシューティングをご覧ください。

次のステップ