建立 TPU 執行個體

本文說明如何建立 TPU 虛擬機器 (VM) 執行個體、連線至該執行個體，以及在該執行個體上執行計算。您在這個快速入門導覽課程中建立的 TPU 執行個體會使用隨選模式，這是預設的消耗選項。

事前準備

  1. 登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎 建立帳戶，親自評估產品在實際工作環境中的成效。新客戶還能獲得價值 $300 美元的免費抵免額，可用於執行、測試及部署工作負載。

  2. 安裝 Google Cloud CLI。

  3. 若您採用的是外部識別資訊提供者 (IdP)，請先使用聯合身分登入 gcloud CLI。

  4. 執行下列指令，初始化 gcloud CLI：

    gcloud init

  5. 建立或選取 Google Cloud 專案。

    選取或建立專案所需的角色

    • 選取專案：選取專案時，不需要具備特定 IAM 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您需要具備專案建立者角色 (roles/resourcemanager.projectCreator)，其中包含 resourcemanager.projects.create 權限。瞭解如何授予角色。

    • 建立 Google Cloud 專案：

      gcloud projects create PROJECT_ID

      將 PROJECT_ID 替換為您要建立的 Google Cloud 專案名稱。

    • 選取您建立的 Google Cloud 專案：

      gcloud config set project PROJECT_ID

      將 PROJECT_ID 替換為 Google Cloud 專案名稱。

  6. 如要使用現有專案進行本指南中的操作，請確認您具有完成本指南所需的權限。如果您建立新專案，則已具備必要權限。

  7. 確認專案已啟用計費功能 Google Cloud 。

  8. 啟用 Compute Engine API：

    啟用 API 時所需的角色

    如要啟用 API，您需要具備服務使用情形管理員 IAM 角色 (roles/serviceusage.serviceUsageAdmin)，其中包含 serviceusage.services.enable 權限。瞭解如何授予角色。

    gcloud services enable compute.googleapis.com

  9. 安裝 Google Cloud CLI。

  10. 若您採用的是外部識別資訊提供者 (IdP)，請先使用聯合身分登入 gcloud CLI。

  11. 執行下列指令，初始化 gcloud CLI：

    gcloud init

  12. 建立或選取 Google Cloud 專案。

    選取或建立專案所需的角色

    • 選取專案：選取專案時，不需要具備特定 IAM 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您需要具備專案建立者角色 (roles/resourcemanager.projectCreator)，其中包含 resourcemanager.projects.create 權限。瞭解如何授予角色。

    • 建立 Google Cloud 專案：

      gcloud projects create PROJECT_ID

      將 PROJECT_ID 替換為您要建立的 Google Cloud 專案名稱。

    • 選取您建立的 Google Cloud 專案：

      gcloud config set project PROJECT_ID

      將 PROJECT_ID 替換為 Google Cloud 專案名稱。

  13. 如要使用現有專案進行本指南中的操作，請確認您具有完成本指南所需的權限。如果您建立新專案，則已具備必要權限。

  14. 確認專案已啟用計費功能 Google Cloud 。

  15. 啟用 Compute Engine API：

    啟用 API 時所需的角色

    如要啟用 API，您需要具備服務使用情形管理員 IAM 角色 (roles/serviceusage.serviceUsageAdmin)，其中包含 serviceusage.services.enable 權限。瞭解如何授予角色。

    gcloud services enable compute.googleapis.com
    16.

必要的角色

如要取得建立 TPU 執行個體、透過 SSH 連線至該執行個體及執行指令所需的權限，請要求管理員授予您專案的下列 IAM 角色：

如要進一步瞭解如何授予角色，請參閱「管理專案、資料夾和組織的存取權」。

您或許也能透過自訂角色或其他預先定義的角色，取得必要權限。

建立 TPU 執行個體

在本地終端機或使用 Cloud Shell 執行下列指令。

  1. 如要建立 TPU 執行個體，請執行下列 gcloud compute instances create 指令：

    gcloud compute instances create quickstart-tpu \
    --machine-type=ct6e-standard-4t \
    --zone=us-east5-b \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accelerator-2204-amd64-with-tpu-v6e

  2. 您也可以使用 gcloud compute instances describe 指令，確認執行個體是否正在執行：

    gcloud compute instances describe quickstart-tpu --zone=us-east5-b

    在輸出內容中，檢查 status 欄位，確認執行個體是否正在執行或仍處於待處理狀態。

連線至 TPU 執行個體

您可以使用 Google Cloud 控制台或 gcloud CLI 連線至 TPU 執行個體。

如要透過 SSH 連線至 VM，請執行下列gcloud compute ssh指令：

gcloud compute ssh quickstart-tpu --project=PROJECT_ID --zone=us-east5-b

安裝 JAX

在 TPU 執行個體上執行下列指令，安裝 JAX 架構：

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

執行系統檢查

在 TPU 執行個體上執行下列指令，確認 JAX 可以存取 TPU 並執行基本作業：

  1. 啟動 Python 3 解譯器：

    python3

  2. 匯入 JAX 模組：

    >>> import jax

  3. 顯示可用的 TPU 核心數：

    >>> jax.device_count()

    輸出內容會顯示 TPU 核心數量。核心數量取決於您使用的 TPU 版本。詳情請參閱「TPU 版本」。

執行計算

  1. 執行下列指令，將兩個數字相加：

    >>> jax.numpy.add(1, 1)

    畫面會顯示如下的輸出內容：

    >>> Array(2, dtype=int32, weak_type=True)

  2. 結束 Python 解譯器：

    >>> exit()

清除所用資源

為了避免系統向您的 Google Cloud 帳戶收取本頁面所用資源的費用，請刪除含有這些資源的 Google Cloud 專案。

如要避免系統向您的 Google Cloud 帳戶收取本頁面所用資源的費用，請使用 gcloud compute instances delete 指令刪除 TPU 執行個體：

gcloud compute instances delete quickstart-tpu --zone=us-east5-b

後續步驟