Cloud TPU マルチスライスの概要

Cloud TPU マルチスライスは、標準のデータ並列処理により、単一のスライス内、または複数の Pod 内のスライスで、トレーニング ジョブが複数の TPU スライスを使用できるようにする、フルスタック パフォーマンス スケーリング テクノロジーです。TPU v4 チップでは、トレーニング ジョブは 1 回の実行で 4,096 個を超えるチップを使用できます。4,096 チップ未満を必要とするトレーニング ジョブの場合、単一スライスが最も高パフォーマンスを発揮します。ただし、複数の小さなスライスの方が簡単に利用できるため、マルチスライスを小さなスライスで使用する場合、起動時間が短縮されます。

複数のスライスによるパフォーマンスの線形スケーリング

マルチスライス構成にデプロイすると、各スライス内の TPU チップがチップ間相互接続（ICI）を介して通信します。異なるスライス内の TPU チップは、CPU（ホスト）にデータを転送することで通信します。CPU は、データセンター ネットワーク（DCN）を介してデータを転送します。マルチスライスでのスケーリングの詳細については、マルチスライスで AI トレーニングを最大数万の Cloud TPU チップまでスケーリングする方法をご覧ください。

マルチスライスのデータフロー

スライス間 DCN 通信を実装するためにデベロッパーがコードを記述することはありません。XLA コンパイラが、そのコードを生成し、最大限のパフォーマンスが発揮できるようにコンピューティングと通信をオーバーラップします。

コンセプト

アクセラレータ タイプ
マルチスライスを構成する各 TPU スライスのシェイプ。マルチスライス リクエスト内の各スライスのアクセラレータ タイプは同じです。アクセラレータ タイプは、TPU タイプ（v4 以降）と TensorCore の数で構成されます。たとえば、v5litepod-128 は、TPU v5e と 128 個の TensorCore を示します。
自動修復
スライスにメンテナンス イベント、プリエンプション、またはハードウェアの障害が発生すると、Cloud TPU が新しいスライスを作成します。新しいスライスを作成するのに十分なリソースがない場合、ハードウェアが利用可能になるまで作成は完了しません。新しいスライスを作成すると、マルチスライス環境内の他のすべてのスライスが再起動され、トレーニングを続行できます。適切に構成された起動スクリプトを使用すると、ユーザーの介入なしに、トレーニング スクリプトが自動的に再起動し、最新のチェックポイントから読み込み、再開します。
データセンター ネットワーキング（DCN）
マルチスライス構成で TPU スライスを接続する、高レイテンシ、低スループット（ICI との比較）のネットワーク。
ギャング スケジューリング
すべての TPU スライスが同時にプロビジョニングされた場合に、すべてのスライスが正常にプロビジョニングされるか、いずれのスライスもプロビジョニングされないことを保証します。
インターチップ相互接続（ICI）
TPU Pod 内で TPU を接続する高速かつ低レイテンシの内部リンク。
マルチスライス
DCN を介して通信できる 2 つ以上の TPU チップスライス。
ノード
マルチスライスのコンテキストでは、ノードは単一の TPU スライスを指します。マルチスライスの各 TPU スライスにはノード ID が割り当てられます。
起動スクリプト
VM が起動または再起動されるたびに実行される標準の Compute Engine 起動スクリプト。マルチスライスの場合、QR 作成リクエストで指定されます。Cloud TPU 起動スクリプトの詳細については、TPU リソースを管理するをご覧ください。
Tensor
ML モデルの多次元データを表すために使用されるデータ構造。
Cloud TPU の容量のタイプ

TPU は、さまざまなタイプの容量から作成できます（TPU の料金の仕組みの使用オプションを参照）。

  • 予約: 予約を使用するには、Google との予約契約が必要です。リソースを作成する際は --reserved フラグを使用します。

  • Spot: Spot VM を使用するプリエンプティブルの割り当てを対象にします。優先度の高いジョブのリクエストに対応できるように、リソースがプリエンプトされる場合があります。リソースを作成する際は --spot フラグを使用します。

  • オンデマンド: 予約を必要とせずプリエンプトされない、オンデマンド割り当てを対象にします。TPU リクエストは、Cloud TPU が提供するオンデマンド割り当てキューに追加されます。リソースの可用性は保証されません。デフォルトで選択されます。フラグは必要ありません。

始める

  1. Cloud TPU 環境を設定します。

  2. In the Google Cloud console, activate Cloud Shell.

    Activate Cloud Shell

    At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

    3. マルチスライスを使用するには、TPU リソースをキューに格納されたリソースとして管理する必要があります。

    入門例

    このチュートリアルでは、MaxText GitHub リポジトリのコードを使用します。MaxText は、Python と Jax で記述された、高パフォーマンスで任意にスケーラブルなオープンソースの十分にテストされた基本 LLM です。Cloud TPU での効率的なトレーニングを目的として設計されています。

    shardings.py のコードは、さまざまな並列化オプションのテストを開始するうえで役立つように設計されています。たとえば、データ並列処理、完全にシャーディングされたデータ並列処理（FSDP）、テンソル並列処理などです。コードは、単一スライス環境からマルチスライス環境にスケーリングされます。

    ICI 並列処理

    ICI は、単一スライスの TPU を接続する高速相互接続を指します。ICI シャーディングは、スライス内のシャーディングに対応します。shardings.py には、次の 3 つの ICI 並列処理パラメータがあります。

    • ici_data_parallelism
    • ici_fsdp_parallelism
    • ici_tensor_parallelism

    これらのパラメータに指定する値によって、各並列化メソッドのシャードの数が決まります。

    これらの入力は、ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism がスライス内のチップの数と等しくなるように制限する必要があります。

    次の表に、v4-8 で使用可能な 4 チップの ICI 並列処理のユーザー入力の例を示します。

    ici_data_parallelism ici_fsdp_parallelism ici_tensor_parallelism
    4 方向 FSDP 1 4 1
    4 方向テンソル並列処理 1 1 4
    2 方向 FSDP + 2 方向テンソル並列処理 1 2 2

    ほとんどの場合、ici_data_parallelism は 1 のままにしておきます。ICI ネットワークは十分高速で、ほぼ常にデータ並列処理よりも FSDP が優先されるためです。

    この例は、JAX を使用して Cloud TPU VM で計算を実行するなど、単一の TPU スライスでのコード実行に精通していることを前提としています。この例は、単一のスライスで shardings.py を実行する方法を示しています。

    1. 環境を設定します。

      $ gcloud auth login
$ export QR_ID=your-queued-resource-id
$ export TPU_NAME=your-tpu-name
$ export PROJECT=your-project-name
$ export ZONE=us-central1-a
$ export NETWORK_NAME=your-network-name
$ export SUBNETWORK_NAME=your-subnetwork-name
$ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
$ export ACCELERATOR_TYPE=v5litepod-16
$ export EXAMPLE_TAG_1=your-tag-1
$ export EXAMPLE_TAG_2=your-tag-2
$ export SLICE_COUNT=4
$ export STARTUP_SCRIPT='#!/bin/bash\n'

      変数の説明

      入力 説明
      QR_ID キューに格納されたリソースのユーザー割り当て ID。
      TPU_NAME ユーザーが割り当てた TPU の名前。
      PROJECT Google Cloud プロジェクト名
      ZONE リソースを作成するゾーンを指定します。
      NETWORK_NAME VPC ネットワークの名前。
      SUBNETWORK_NAME VPC ネットワーク内のサブネットの名前。
      RUNTIME_VERSION Cloud TPU ソフトウェアのバージョン。
      ACCELERATOR_TYPE v4-16
      EXAMPLE_TAG_1、EXAMPLE_TAG_2 … ネットワーク ファイアウォールの有効なソースやターゲットを識別するために使用されるタグ。
      SLICE_COUNT スライスの数。上限は 256 スライスです。
      STARTUP_SCRIPT 起動スクリプトを指定すると、TPU スライスがプロビジョニングまたは再起動されたときにスクリプトが実行されます。

    2. gcloud の SSH 認証鍵を作成します。パスワードは空白のままにすることをおすすめします（次のコマンドの実行後に 2 回 Enter を押します）。google_compute_engine ファイルがすでに存在しているというメッセージが表示された場合は、既存のバージョンを置き換えます。

      $ ssh-keygen -f ~/.ssh/google_compute_engine

    3. TPU をプロビジョニングします。

      gcloud

      $ gcloud compute tpus queued-resources \
    create ${QR_ID} \
    --accelerator-type=${ACCELERATOR_TYPE} \
    --runtime-version=${RUNTIME_VERSION} \
    --node-id=${TPU_NAME} \
    --zone=${ZONE} \
    [--reserved |--spot]

      Google Cloud CLI では、タグなどの QR コードの作成オプションはサポートされていません。詳細については、QR を作成するをご覧ください。

      コンソール

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