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

このドキュメントでは、Tensor Processing Unit（TPU）を搭載した仮想マシン（VM）インスタンスを使用する Batch ジョブを作成して実行する方法について説明します。TPU は、機械学習（ML）と人工知能（AI）のワークロードを高速化するために設計された、Google 独自に開発された特定用途向け集積回路（ASIC）です。Batch ジョブを作成するときに、必要に応じて TPU マシンタイプを指定して、コンテナ化された AI / ML 推論、評価、ファインチューニング、トレーニング ワークロード、またはスクリプトベースの AI / ML 推論、評価、ファインチューニング、トレーニング ワークロードを高速化できます。

TPU の詳細については、 Google Cloudの TPU の概要と、Compute Engine ドキュメントの TPU マシンタイプをご覧ください。または、代わりに GPU を使用して、データ処理、AI、ML ワークロードを高速化することもできます。

始める前に

  1. Batch を以前に使用したことがない場合は、Batch を使ってみるを確認し、プロジェクトとユーザーの前提条件を完了して Batch を有効にします。

  2. ジョブの作成に必要な権限を取得するには、次の IAM ロールを付与するよう管理者に依頼してください。

    ロールの付与については、プロジェクト、フォルダ、組織へのアクセスを管理するをご覧ください。

    必要な権限は、カスタムロールや他の事前定義ロールから取得することもできます。

制限事項

ジョブで TPU を使用する場合は、TPU に関する Compute Engine の制限事項に加えて、次の制限事項がすべて適用されます。

  • 単一ホスト TPU スライスのみ: Batch は、単一ホスト TPU スライスと呼ばれる独立した TPU VM のみをサポートします。相互接続された複数の VM で構成されるマルチホスト TPU スライスはサポートされていません。

    ジョブは、1 つ以上の単一ホスト TPU スライスで実行できます。たとえば、ジョブで taskCount が 4 に、taskCountPerNode が 1 に設定されている場合、ジョブが実行される VM の数は 4 以下になります。マルチホスト TPU スライスとは異なり、単一ホスト TPU スライスはそれぞれ独立した VM として動作し、高帯域幅、低レイテンシの通信用のチップ間相互接続（ICI）ネットワーキングを提供しません。

    単一ホスト TPU スライスとマルチホスト TPU スライスの詳細については、Compute Engine のドキュメントをご覧ください。

  • サポートされている TPU バージョンとマシンタイプ: Batch は、次の TPU バージョンとマシンタイプの単一ホスト TPU スライス（独立した TPU VM）のみをサポートしています。

    • TPU7x（Ironwood）: tpu7x-standard-4t

    • v6e（Trillium）: ct6e-standard-1t、ct6e-standard-4t、ct6e-standard-8t

    • v5p: ct5p-hightpu-4t

    サポートされているディスクなど、各 TPU マシンタイプの詳細については、Compute Engine ドキュメントの TPU マシンタイプをご覧ください。

  • サポートされている消費オプション: Batch は、TPU を使用するジョブで Flex Start VM（FLEX_START プロビジョニング モデル）をサポートしていません。Flex Start VM を除き、Compute Engine ドキュメントの TPU バージョン別の使用オプションの可用性に従って、ジョブの TPU バージョンでサポートされている使用オプションを指定できます。

  • 単一ゾーンのロケーションが必要: TPU を使用するジョブを作成する場合は、ジョブの許可されたロケーション（allowedLocations[] フィールド）を、TPU バージョンまたは予約（該当する場合）を含むゾーンに設定して、ジョブを特定のゾーンに制限する必要があります。

  • サポートされていない機能: TPU を使用するジョブは、次の機能をサポートしていません。

TPU を使用する Job を作成する

TPU を使用するジョブを作成するには、次の操作を行います。

  1. TPU を使用するジョブの構成を計画する。
  2. 特定した構成を使用してジョブを作成します。TPU を使用するジョブを作成する方法の例については、このドキュメントの TPU を使用するジョブを作成するをご覧ください。

TPU を使用するジョブの構成を計画する

TPU を使用するジョブを作成する前に、次の操作を行う必要があります。

  1. TPU のバージョンを選択する: ワークロードに適した TPU のバージョンを選択します。ワークロード タイプ別の TPU バージョンの一覧については、ワークロード タイプ別の推奨 TPU バージョンをご覧ください。

    サポートされている TPU バージョンのいずれかを選択してください。

  2. TPU の使用オプションを選択する: ワークロード、その期間、費用のニーズに最適な使用オプションを選択します。サポートされている使用オプションのいずれかを選択してください。次の表に、各使用オプションのユースケースとジョブ要件をまとめます。

    使用オプションとジョブ要件

    予約にバインド（カレンダー モードの将来の予約または容量ブロックの将来の予約）

    • ユースケース: 特定の期間に TPU リソースの可用性を非常に高いレベルで保証する必要があるジョブには、予約にバインドされた使用オプションをおすすめします。Compute Engine では、これらはカレンダー モードの将来の予約（1 ～ 90 日間）と容量ブロックの将来の予約（1 年以上）として記述されます。

    • 詳細: ジョブを作成する前に、選択した TPU バージョンと RESERVATION_BOUND プロビジョニング モデルの予約を取得または特定する必要があります。予約を作成すると、その予約を削除するまで、指定された VM を実行するのと同じ料金で VM の費用が発生します。予約を使用している VM に対しては別途料金は発生しませんが、予約には使用量に関係なく料金が発生します。

      詳細については、Compute Engine ドキュメントの予約タイプを選択するをご覧ください。

    Batch は、次のすべての処理を行うジョブに予約バインド予約を使用します。

    • provisioningModel フィールドを RESERVATION_BOUND に設定します。
    • 互換性のある予約を明示的に対象とする。reservation フィールドを、予約、予約ブロック、予約サブブロックの略称（予約が同じプロジェクトとゾーンにある場合）または完全なリソースパスに設定します。

    オンデマンド

    • ユースケース: オンデマンドは高速で簡単に試用できるため、ほとんどのジョブのデフォルトの消費オプションです。ただし、別の使用オプションが優先事項に合っている場合や、リソースの取得可能性に関する問題が繰り返し発生する場合は、オンデマンドをおすすめしません。

    • 詳細: 通常、オンデマンドは Compute Engine VM にアクセスするデフォルトの方法です。オンデマンドでは、リクエスト時にリソースにアクセスできます（利用可能な場合）。リソースは、ジョブが完了するか、実行時間の上限に達するまで削除されません。

    Batch は、デフォルトでほとんどのジョブにオンデマンドを使用します。ジョブでオンデマンドを使用する場合は、次のすべてを行います。

    • provisioningModel フィールドを STANDARD に設定します。
    • 予約済みリソースの消費を防止する。ジョブが互換性のある予約を自動的に使用しないようにするには、reservation フィールドを NO_RESERVATION に設定します。詳細については、予約済み VM を消費できないジョブを作成して実行するをご覧ください。

    Spot VM

    • ユースケース: Spot VM を使用して、フォールト トレラントなワークロードの費用を削減することをおすすめします。

    • 詳細: Spot VM を使用するジョブの作成リクエストが処理されると、VM には Spot VM の料金が適用されます。この料金は頻繁に変更される可能性があり、オンデマンド料金から最大 91% の割引が適用される場合があります。ただし、Spot VM は常に利用可能であるとは限らず、いつでもプリエンプトできます。

      Spot VM と Spot VM を使用するジョブの推奨事項の詳細を確認する。

    Batch は、次のすべての処理を行うジョブに Spot VM を使用します。

    • provisioningModel フィールドを SPOT に設定します。
    • （省略可）予約の使用を防止します。予約で Spot VM を使用することはできません。具体的には、ジョブで reservation フィールドを NO_RESERVATION に設定できます。詳細については、予約済み VM を消費できないジョブを作成して実行するをご覧ください。

  3. 選択した TPU が利用可能なゾーンを選択する: TPU バージョンを使用するには、ジョブの allowedLocations[] フィールドを、選択した TPU を含むゾーンに設定する必要があります。

    • 予約を使用する消費オプションを選択した場合は、予約が配置されているゾーンを使用する必要があります。
    • それ以外の場合は、TPU の可用性に従って、TPU バージョンを含むゾーンを選択します。

    どちらの場合も、ゾーンは Batch でサポート対象の地域のいずれかに存在する必要があります。

  4. プロジェクトに十分な TPU 割り当てがあることを確認する: オンデマンド VM または Spot VM で TPU を使用するジョブを作成する場合は、使用するリージョンに十分な TPU 割り当てが必要です。TPU 予約を使用する TPU を使用するジョブを作成する場合、予約の作成時に割り当てが使用されるため、TPU の割り当ては必要ありません。TPU 割り当て名の一覧については、TPU 割り当てをご覧ください。割り当てを表示する方法については、割り当ての表示と管理をご覧ください。

  5. TPU マシンタイプを選択する: 選択した TPU バージョンの TPU マシンタイプを選択します。予約を使用する使用オプションを選択した場合は、予約のマシンタイプを使用する必要があります。各 TPU バージョンで使用可能なマシンタイプのリストについては、TPU マシンタイプをご覧ください。

    制限事項により、TPU マシンタイプには単一ホストのトポロジを使用する必要があります。

TPU を使用する Job を作成する

以降のセクションでは、各消費オプションで TPU を使用するサンプルジョブを作成する方法について説明します。

使用オプション ジョブ指示の例
予約にバインド（カレンダー モードの将来の予約または容量ブロックの将来の予約）

予約に制限された使用オプションで TPU を使用する
オンデマンドまたは Spot VM

オンデマンド VM または Spot VM で TPU を使用する

オンデマンド VM または Spot VM で TPU を使用する

gcloud CLI または Batch API を使用して、オンデマンド VM または Spot VM で TPU を使用するジョブを作成できます。

gcloud

  1. サポートされている単一ホスト TPU マシンタイプを指定し、プロビジョニング モデルを STANDARD または SPOT に設定し、TPU マシンタイプを提供するロケーションで実行する JSON 構成ファイルを作成します。

    たとえば、基本的なスクリプトのジョブを作成するには、次の内容の JSON ファイルを作成します。

    {
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "PROVISIONING_MODEL",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

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

    • PROVISIONING_MODEL: 使用オプションのプロビジョニング モデル。オンデマンドの場合は STANDARD、Spot VM の場合は SPOT。

    • MACHINE_TYPE: サポートされている単一ホストの TPU マシンタイプ（例: ct5p-hightpu-4t）。

    • ALLOWED_LOCATIONS: ジョブに許可されるロケーション。TPU を使用するジョブの場合、allowedLocations[] フィールドを、TPU マシンタイプを提供するゾーンに設定する必要があります。たとえば、zones/us-central1-a は TPU マシンタイプ ct5p-hightpu-4t を提供します。

  2. ジョブを作成して実行するには、gcloud batch jobs submit コマンドを使用します:

    gcloud batch jobs submit JOB_NAME \
    --location LOCATION \
    --config JSON_CONFIGURATION_FILE

    以下を置き換えます。

    • JOB_NAME: ジョブの名前。

    • LOCATION: ジョブのロケーション。

    • JSON_CONFIGURATION_FILE: ジョブの構成の詳細を含む JSON ファイルへのパス。

API

サポートされている単一ホスト TPU マシンタイプを指定し、プロビジョニング モデルを STANDARD または SPOT に設定し、TPU マシンタイプを提供するロケーションで実行する jobs.create メソッドに POST リクエストを送信します。

たとえば、基本的なスクリプト ジョブを作成するには、次のリクエストを行います。

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "PROVISIONING_MODEL",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

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

  • PROJECT_ID: プロジェクトのプロジェクト ID

  • LOCATION: ジョブのロケーション。

  • JOB_NAME: ジョブの名前。

  • PROVISIONING_MODEL: 使用オプションのプロビジョニング モデル。オンデマンドの場合は STANDARD、Spot VM の場合は SPOT。

  • MACHINE_TYPE: サポートされている単一ホストの TPU マシンタイプ（例: ct5p-hightpu-4t）。

  • ALLOWED_LOCATIONS: ジョブに許可されるロケーション。TPU を使用するジョブの場合、allowedLocations[] フィールドを、TPU マシンタイプを提供するゾーンに設定する必要があります。たとえば、zones/us-central1-a は TPU マシンタイプ ct5p-hightpu-4t を提供します。

予約に制限された使用オプションで TPU を使用する

gcloud CLI または Batch API を使用して、予約バインド（RESERVATION_BOUND）プロビジョニング モデルの消費オプション（容量ブロックの将来の予約またはカレンダー モードの将来の予約）で TPU を使用するジョブを作成できます。

gcloud

  1. 単一ホスト TPU マシンタイプの予約を指定し、プロビジョニング モデルを RESERVATION_BOUND に設定し、予約のゾーンで実行する JSON 構成ファイルを作成します。

    たとえば、基本的なスクリプトのジョブを作成するには、次の内容の JSON ファイルを作成します。

    {
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "RESERVATION_BOUND",
          "reservation": "RESERVATION",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

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

    • RESERVATION: 予約への次のいずれかのパスに置き換えます。

      • カレンダー モードの将来の予約、または容量ブロックの将来の予約の一部を使用する場合は、予約を指定します。

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME

      • 容量ブロックの将来の予約の特定のブロックを使用する場合は、予約ブロックを指定します。

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME

      • 容量ブロックの将来の予約の特定のサブブロックを使用する場合は、予約サブブロックを指定します。

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME

    • MACHINE_TYPE: 予約の TPU マシンタイプ（例: ct5p-hightpu-4t）。

    • ALLOWED_LOCATIONS: ジョブに許可されるロケーション。予約済みの TPU を使用するジョブの場合、allowedLocations[] フィールドを予約のゾーンに設定する必要があります。たとえば、zones/us-central1-a は TPU マシンタイプ ct5p-hightpu-4t を提供します。

  2. ジョブを作成して実行するには、gcloud batch jobs submit コマンドを使用します:

    gcloud batch jobs submit JOB_NAME \
    --location LOCATION \
    --config JSON_CONFIGURATION_FILE

    以下を置き換えます。

    • JOB_NAME: ジョブの名前。

    • LOCATION: ジョブのロケーション。

    • JSON_CONFIGURATION_FILE: ジョブの構成の詳細を含む JSON ファイルへのパス。

API

単一ホスト TPU マシンタイプの予約を指定し、プロビジョニング モデルを RESERVATION_BOUND に設定し、予約のゾーンで実行する jobs.create メソッドに POST リクエストを送信します。

たとえば、基本的なスクリプト ジョブを作成するには、次のリクエストを行います。

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
  "taskGroups": [
    {
      "taskSpec": {
        "runnables": [
          {
            "script": {
              "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "RESERVATION_BOUND",
          "reservation": "RESERVATION",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

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

  • PROJECT_ID: プロジェクトのプロジェクト ID

  • LOCATION: ジョブのロケーション。

  • JOB_NAME: ジョブの名前。

  • RESERVATION: 予約への次のいずれかのパスに置き換えます。

    • カレンダー モードの将来の予約、または容量ブロックの将来の予約の一部を使用する場合は、予約を指定します。

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME

    • 容量ブロックの将来の予約の特定のブロックを使用する場合は、予約ブロックを指定します。

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME

    • 容量ブロックの将来の予約の特定のサブブロックを使用する場合は、予約サブブロックを指定します。

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME

  • MACHINE_TYPE: 予約の TPU マシンタイプ（例: ct5p-hightpu-4t）。

  • ALLOWED_LOCATIONS: ジョブに許可されるロケーション。予約済みの TPU を使用するジョブの場合、allowedLocations[] フィールドを予約のゾーンに設定する必要があります。たとえば、zones/us-central1-a は TPU マシンタイプ ct5p-hightpu-4t を提供します。

予約を使用するジョブの作成の詳細については、VM 予約を使用してリソースの可用性を確保するをご覧ください。

次のステップ