TPU を使用したモデルのトレーニング

Tensor Processing Unit(TPU)は、機械学習ワークロードを高速化するために使用される、Google 独自開発の ASIC です。Cloud TPU を使用して、AI Platform Training でトレーニング ジョブを実行できます。AI Platform Training にはジョブ管理インターフェースが用意されているので、TPU を自分で管理する必要はありません。CPU や GPU でトレーニングを実行する場合と同じように AI Platform Training jobs API を使用できます。

高レベルの TensorFlow API は、Cloud TPU ハードウェア上でモデルを実行する場合に役立ちます。

Google Cloud 環境の設定

スタートガイドのセットアップ セクションに沿って、Google Cloud 環境を構成します。

Cloud TPU にプロジェクトへのアクセスを許可する

Google Cloud プロジェクトに関連付けられた Cloud TPU サービス アカウント名を承認する手順は次のとおりです。

  1. projects.getConfig を呼び出して、Cloud TPU サービス アカウント名を取得します。例:

    PROJECT_ID=PROJECT_ID
    
    curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
        https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
    
  2. API から返された serviceAccountProjecttpuServiceAccount フィールドの値を保存します。

  3. Cloud TPU サービス アカウントを初期化します。

    curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
      -H "Content-Type: application/json" -d '{}'  \
      https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
    

Cloud ML サービス エージェントのロールを使用して、Cloud TPU サービス アカウントをプロジェクト内のメンバーとして追加します。Google Cloud Console または gcloud コマンドを使用して、次の手順を行います。

コンソール

  1. Google Cloud Console にログインし、TPU を使用しているプロジェクトを選択します。
  2. [IAM と管理] > [IAM] を選択します。
  3. [追加] ボタンをクリックして、プロジェクトにメンバーを追加します。
  4. [メンバー] テキスト ボックスに TPU サービス アカウントを入力します。
  5. [ロール] プルダウン リストをクリックします。
  6. Cloud ML サービス エージェントのロールを有効にします([サービス エージェント] > [Cloud ML サービス エージェント])。

gcloud

  1. プロジェクト ID と Cloud TPU サービス アカウントが格納されている環境変数を設定します。

    PROJECT_ID=PROJECT_ID
    SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
    
  2. Cloud TPU サービス アカウントに ml.serviceAgent ロールを付与します。

    gcloud projects add-iam-policy-binding $PROJECT_ID \
        --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
    

サービス アカウントへのロールの付与については、IAM のドキュメントをご覧ください。

例: サンプル MNIST モデルのトレーニング

このセクションでは、TPU とランタイム バージョン 2.11 を使用してサンプル MNIST モデルをトレーニングする方法について説明します。この例のジョブでは、マシン構成にあらかじめ定義された BASIC_TPU スケールティアを使用します。ガイドのこの後のセクションでは、カスタム構成を設定する方法について説明します。

この例では、gcloud CLI がインストールされた Bash シェルを使用していることを前提としています。次のコマンドを実行してコードを取得し、トレーニング ジョブを AI Platform Training に送信します。

  1. TensorFlow のリファレンス モデルのコードをダウンロードし、サンプルコードのあるディレクトリに移動します。

    git clone https://github.com/tensorflow/models.git \
      --branch=v2.11.0 \
      --depth=1
    
    cd models
    
  2. models ディレクトリに setup.py ファイルを作成します。これにより、gcloud ai-platform jobs submit training コマンドは、トレーニング コードの tarball を作成するときに、必要なすべてのサブパッケージを models/official ディレクトリに含め、トレーニング ジョブ実行時に AI Platform Training が TensorFlow データセットを依存関係としてインストールします。このトレーニング コードは、MNIST データを読み込むために TensorFlow データセットに依存します。

    setup.py ファイルを作成するには、シェルで次のコマンドを実行します。

    cat << END > setup.py
    from setuptools import find_packages
    from setuptools import setup
    
    setup(
        name='official',
        install_requires=[
           'tensorflow-datasets~=3.1',
           'tensorflow-model-optimization>=0.4.1'
       ],
        packages=find_packages()
    )
    END
    
  3. gcloud ai-platform jobs submit training コマンドを使用して、トレーニング ジョブを送信します。

    gcloud ai-platform jobs submit training tpu_mnist_1 \
      --staging-bucket=gs://BUCKET_NAME \
      --package-path=official \
      --module-name=official.vision.image_classification.mnist_main \
      --runtime-version=2.11 \
      --python-version=3.7 \
      --scale-tier=BASIC_TPU \
      --region=us-central1 \
      -- \
      --distribution_strategy=tpu \
      --data_dir=gs://tfds-data/datasets \
      --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
    

    BUCKET_NAME は、Google Cloud プロジェクト内の Cloud Storage バケットの名前に置き換えます。gcloud CLI はパッケージ化されたトレーニング コードをこのバケットにアップロードし、AI Platform Training はバケットにトレーニング出力を保存します。

  4. トレーニング ジョブをモニタリングします。ジョブが完了したら、gs://BUCKET_NAME/tpu_mnist_1_output ディレクトリでその出力を表示できます。

Cloud TPU でのモデルのトレーニングの詳細

このセクションでは、AI Platform Training 上で Cloud TPU を使用してジョブの構成とモデルのトレーニングを行う方法について詳しく説明します。

TPU を提供するリージョンの指定

TPU が利用可能なリージョンでジョブを実行する必要があります。次のリージョンで現在 TPU にアクセスできるようになっています。

  • us-central1
  • europe-west4

モデル トレーニングやオンライン / バッチ予測などの AI Platform Training サービスで利用可能なリージョンの詳細については、リージョンのガイドをご覧ください。

TensorFlow と AI Platform Training のバージョニング

AI Platform Training ランタイム バージョン 1.15、2.1、2.2、2.3、2.4、2.5、2.6、2.7、2.8、2.9、2.11 は Cloud TPU でのモデルのトレーニングに使用できます。AI Platform Training ランタイム バージョンと対応する TensorFlow バージョンに関する詳細をご覧ください。

バージョニング ポリシーは Cloud TPU と同じです。トレーニングのジョブ リクエストでは、TPU で使用可能なランタイム バージョンを指定し、トレーニング コードで使用されている TensorFlow バージョンと一致することを確認してください。

TPU gRPC サーバーへの接続

TensorFlow プログラムでは、TPUClusterResolver を使用して TPU VM 上で実行されている TPU gRPC サーバーに接続します。

TPU の使用に関する TensorFlow ガイドでは、TPUStrategy 配信戦略で TPUClusterResolver を使用する方法を説明しています。

ただし、AI Platform Training で実行するコードに TPUClusterResolver を使用する場合は、重要な変更点があります。TPUClusterResolver インスタンスの作成時に引数を指定しないでください。tpuzoneproject の各キーワード引数がすべてデフォルト値の None に設定されている場合、AI Platform Training は環境変数を通じて必要な接続の詳細を含むクラスタ リゾルバを自動的に提供します。

次の TensorFlow 2 の例は、AI Platform Training でトレーニングするためのクラスタ リゾルバと分散戦略を初期化する方法を示しています。

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

TensorFlow コードでの TPU の使用

マシンで TPU を使用するには、TensorFlow 2 の TPUStrategy API を使用します。これを行う方法については、TPU の使用に関する TensorFlow ガイドをご覧ください。

TensorFlow 1 で TPU を使用してトレーニングするには、代わりに TPUEstimator API を使用します。これを行う方法については、TPUEstimator API に関する Cloud TPU ガイドをご覧ください。

Cloud TPU のドキュメントには、Cloud TPU で利用できる簡単な TensorFlow 演算のリストも記載されています。

PyTorch コードでの TPU の使用

ビルド済みの PyTorch コンテナを使用する際に TPU を使用するには、torch_xla パッケージを使用します。PyTorch ドキュメントのトレーニングで TPU に torch_xla を使用する方法について学習します。torch_xla の他の使用例については、PyTorch XLA GitHub リポジトリのチュートリアルをご覧ください。

AI Platform Training で TPU を使用してトレーニングする場合は、複数の XLA デバイスではなく、単一の XLA デバイスを使用します。

このページの次のセクションで、PyTorch と TPU のトレーニング ジョブの構成もご覧ください。

カスタム TPU マシンの構成

TPU トレーニング ジョブは、2 つの VM による構成で実行されます。一方の VM(マスター)が Python コードを実行します。マスターは、TPU ワーカーで実行されている TensorFlow サーバーを駆動します。

AI Platform Training で TPU を使用するには、次の 3 つの方法のいずれかで TPU 対応マシンにアクセスするようにトレーニング ジョブを構成します。

  • BASIC_TPU スケール階層を使用する。この方法で TPU v2 アクセラレータにアクセスできます。
  • マスター VM で cloud_tpu ワーカーと従来のマシンタイプを使用する。この方法で TPU v2 アクセラレータにアクセスできます。
  • マスター VM で cloud_tpu ワーカーと Compute Engine マシンタイプを使用する。この方法で TPU v2 アクセラレータまたは TPU v3 アクセラレータにアクセスできます。TPU v3 アクセラレータはベータ版で利用できます。

基本的な TPU 対応マシン

スケールティアを BASIC_TPU に設定して、前述のサンプルを実行した場合と同様に、マスター VM と、8 つの TPU v2 コアを搭載した 1 つの TPU を使用している TPU VM を取得します。

従来のマシンタイプ構成の TPU ワーカー

また、マスター VM でさらに多くのコンピューティング リソースが必要な場合は、次のようにカスタムマシン構成を設定することもできます。

  • スケールティアを CUSTOM に設定する。
  • ジョブの要件に合った従来のマシンタイプを使用するようにマスター VM を構成する。
  • workerTypecloud_tpu に設定して、8 つの TPU v2 コアを搭載した 1 つの Cloud TPU を使用している TPU VM を取得します。
  • workerCount を 1 に設定する。
  • Cloud TPU を使用する場合は、パラメータ サーバーを指定しない。parameterServerCount がゼロより大きい場合、サービスはジョブ リクエストを拒否します。

次の例は、このタイプの構成を使用する config.yaml ファイルを示しています。

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

Compute Engine マシンタイプ構成の TPU ワーカー

マスター VM の Compute Engine マシンタイプにカスタムマシン構成を設定し、TPU VM に接続された acceleratorConfig を設定することもできます。

このタイプの構成を使用して、8 つの TPU v2 コアを搭載した TPU ワーカー(acceleratorConfig なしの構成に類似)または 8 つの v3 コアを搭載した TPU ワーカー(ベータ版)を設定できます。TPU v2 アクセラレータと TPU v3 アクセラレータの違いをご確認ください。

Compute Engine マシンタイプを使用すると、次のようにマスター VM をより柔軟に構成できます。

  • スケールティアを CUSTOM に設定する。
  • ジョブ要件に合った Compute Engine マシンタイプを使用するようにマスター VM を構成する。
  • workerTypecloud_tpu に設定する。
  • acceleratorConfig フィールドを含む workerConfig を追加します。その acceleratorConfig 内で、typeTPU_V2 または TPU_V3 に、count8 に設定します。それ以外の数の TPU コアを接続することはできません。
  • workerCount を 1 に設定する。
  • Cloud TPU を使用する場合は、パラメータ サーバーを指定しない。parameterServerCount がゼロより大きい場合、サービスはジョブ リクエストを拒否します。

次の例は、このタイプの構成を使用する config.yaml ファイルを示しています。

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3(ベータ版)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3
      count: 8

TPU Pod の使用

TPU Pod は、専用の高速ネットワーク インターフェースによって接続された TPU デバイスのコレクションです。1 つの TPU Pod に最大 2,048 個の TPU コアを設定して、複数の TPU で処理負荷を分散できます。

TPU Pod を使用するには、まず割り当ての増加リクエストを送信する必要があります。

次の config.yaml ファイルの例は、TPU Pod の使用方法を示しています。

TPU v2 Pod

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2_POD
      count: 128

TPU v3 Pod

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3_POD
      count: 32

TPU タイプごとに使用できる Pod コアの数には制限があります。有効な構成:

TPU Pod タイプ 使用可能な Pod のコア数
TPU_V2_POD 32, 128, 256, 512
TPU_V3_POD 32, 128, 256

TPU Pod コアを最大限に活用する方法の詳細については、TPU Pod に関する Cloud TPU のドキュメントをご覧ください。

TPU ワーカーでのビルド済み PyTorch コンテナの使用

TPU を使用して PyTorch トレーニングを実行する場合、トレーニング ジョブの trainingInputtpuTfVersion フィールドを指定する必要があります。トレーニングに使用するビルド済みの PyTorch コンテナのバージョンに合わせて tpuTfVersion を設定します。

AI Platform Training は、次のビルド済み PyTorch コンテナの TPU によるトレーニングをサポートしています。

コンテナ イメージの URI tpuTfVersion
gcr.io/cloud-ml-public/training/pytorch-xla.1-11 pytorch-1.11
gcr.io/cloud-ml-public/training/pytorch-xla.1-10 pytorch-1.10
gcr.io/cloud-ml-public/training/pytorch-xla.1-9 pytorch-1.9
gcr.io/cloud-ml-public/training/pytorch-xla.1-7 pytorch-1.7
gcr.io/cloud-ml-public/training/pytorch-xla.1-6 pytorch-1.6

たとえば、PyTorch 1.11 のビルド済みコンテナを使用してトレーニングするには、次の config.yaml ファイルを使用してトレーニングを構成します。

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
    tpuTfVersion: pytorch-1.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

このページの前のセクションの PyTorch コードでの TPU の使用もご覧ください。

TPU ワーカーでのカスタム コンテナの使用

TPU をサポートする AI Platform Training ランタイム バージョンのいずれかを使用せずに、TPU ワーカーでカスタム コンテナを実行する場合は、トレーニング ジョブの送信時に追加の構成フィールドを指定する必要があります。tpuTfVersion を、コンテナで使用する TensorFlow のバージョンが含まれているランタイム バージョンに設定します。TPU でのトレーニングのために現在サポートされているランタイム バージョンを指定する必要があります。

カスタム コンテナを使用するジョブを構成しているので、AI Platform のトレーニングではトレーニング ジョブの実行時にこのランタイム バージョンの環境を使用しません。ただし、カスタム コンテナで使用する TensorFlow のバージョン用に TPU ワーカーを適切に準備できるように、AI Platform のトレーニングでこのフィールドが必要です。

次の例では、前のセクションと同様の TPU 構成の config.yaml ファイルを示しますが、この場合はマスター VM と TPU ワーカーが異なるカスタム ワーカーで動作します。

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3(ベータ版)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V3
      count: 8

トレーニング ジョブを送信するために gcloud beta ai-platform jobs submit training コマンドを使用する場合は、config.yaml ファイルではなく --tpu-tf-version フラグを使用して tpuTfVersion API フィールドを指定できます。

TPU のプロビジョニング後の TPUClusterResolver の使用

カスタム コンテナを使用する場合は、TPUClusterResolver を呼び出して使用する前に TPU がプロビジョニングされるのを待つ必要があります。次のサンプルコードは、TPUClusterResolver ロジックの処理方法を示しています。

def wait_for_tpu_cluster_resolver_ready():
  """Waits for `TPUClusterResolver` to be ready and return it.

  Returns:
    A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
    return None.
  Raises:
    RuntimeError: if failed to schedule TPU.
  """
  tpu_config_env = os.environ.get('TPU_CONFIG')
  if not tpu_config_env:
    tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
    return None

  tpu_node = json.loads(tpu_config_env)
  tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)

  num_retries = 40
  for i in range(num_retries):
    try:
      tpu_cluster_resolver = (
          tf.contrib.cluster_resolver.TPUClusterResolver(
              tpu=[tpu_node['tpu_node_name']],
              zone=tpu_node['zone'],
              project=tpu_node['project'],
              job_name='worker'))
      tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
      if 'worker' in tpu_cluster_resolver_dict:
        tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
        return tpu_cluster_resolver
    except Exception as e:
      if i < num_retries - 1:
        tf.logging.info('Still waiting for provisioning of TPU VM instance.')
      else:
        # Preserves the traceback.
        raise RuntimeError('Failed to schedule TPU: {}'.format(e))
    time.sleep(10)

  # Raise error when failed to get TPUClusterResolver after retry.
  raise RuntimeError('Failed to schedule TPU.')

詳細については、カスタム コンテナを使用した分散トレーニングをご覧ください。

次のステップ