建立多主機 TPU 配量

建立多主機 TPU 配量

瞭解如何使用代管執行個體群組 (MIG) 建立多主機 TPU 配量、連線至配量，以及執行計算。本快速入門導覽課程會使用隨選方案。在本機終端機或 Cloud Shell 中執行本快速入門導覽課程的指令。

事前準備

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

安裝 Google Cloud CLI。

注意：如果您先前已安裝 gcloud CLI，請執行 gcloud components update，確認您使用的是最新版本。

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

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

gcloud init

建立或選取 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 專案名稱。

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

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

啟用 Compute Engine API：

啟用 API 時所需的角色

您必須具備 serviceusage.services.enable 權限，才能啟用 API。如果您建立了專案，可能已透過「擁有者」角色 (roles/owner) 取得這項權限。否則，您可以透過「服務使用情形管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。

gcloud services enable compute.googleapis.com

安裝 Google Cloud CLI。

注意：如果您先前已安裝 gcloud CLI，請執行 gcloud components update，確認您使用的是最新版本。

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

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

gcloud init

建立或選取 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 專案名稱。

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

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

啟用 Compute Engine API：

啟用 API 時所需的角色

您必須具備 serviceusage.services.enable 權限，才能啟用 API。如果您建立了專案，可能已透過「擁有者」角色 (roles/owner) 取得這項權限。否則，您可以透過「服務使用情形管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。

gcloud services enable compute.googleapis.com

必要的角色

如要取得建立 MIG (可形成多主機 TPU 節點) 所需的權限、使用 SSH 連線至 MIG 中的每個 VM，以及執行指令，請要求管理員在專案中授予您下列 IAM 角色：

• Compute 執行個體管理員 (v1) (roles/compute.instanceAdmin.v1)
• 服務帳戶使用者 (roles/iam.serviceAccountUser)
• 服務使用情形管理員 (roles/serviceusage.serviceUsageAdmin)

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

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

建立執行個體範本

如要為 TPU v6e VM 建立執行個體範本，請使用 gcloud compute instance-templates create 指令：

注意： 如果使用其他機器類型，請務必使用支援的作業系統映像檔。如需各個支援的 TPU 版本適用的 Google 管理 OS 映像檔清單，請參閱「TPU OS 映像檔」一文。

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

建立工作負載政策

工作負載政策會定義運算執行個體的實體屬性。在 TPU 晶片中，加速器拓撲會定義 TPU 晶片的實體排列方式。互連多主機 TPU 節點必須指定加速器拓撲。

如要為多主機 TPU 配量建立工作負載政策，請使用 gcloud compute resource-policies create workload-policy 指令搭配 --accelerator-topology 旗標。下列指令會建立具有 2x4 拓撲的工作負載政策：

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

建立 MIG

執行下列指令，建立構成多主機 TPU 節點的 MIG。

注意： MIG 排隊等候容量時，您可能會遇到暫時的資源可用性錯誤。如果等待時間過長，可以嘗試其他區域。如需支援的地點清單，請參閱 TPU 區域和可用區。如果選擇其他區域，請先重複「建立工作負載政策」步驟，在該區域建立工作負載政策。請務必在後續所有指令中更新區域。

1. 如要建立構成多主機 TPU 配量的 MIG，請使用 gcloud compute instance-groups managed create 指令：

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   將 PROJECT_ID 替換為 Google Cloud 專案 ID。

2. (選用) 使用下列指令，確認代管執行個體正在執行：

   • 如要查看整體 MIG 狀態，請使用 gcloud compute instance-groups managed describe 指令：

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • 如要查看每個執行個體的狀態，請使用 gcloud compute instance-groups managed list-instances 指令：

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

安裝 JAX

在 MIG 中所有 TPU VM 執行個體的虛擬環境中，安裝依附元件和 JAX 架構。如果 TPU VM 安裝的 Python 版本低於 3.11，您必須安裝 Python 3.11，才能執行最新版 JAX。

1. 檢查 TPU VM 上執行的 Python 版本：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   如果版本低於 Python 3.11，請安裝 Python 3.11：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. 建立虛擬環境：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. 在虛擬環境中安裝 JAX：

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

在配量上執行 JAX 程式碼

如要在 TPU 配量上執行 JAX 程式碼，您必須在 TPU 配量中的每個主機上執行程式碼。在對切片中的每個主機呼叫 jax.device_count() 函式之前，該函式呼叫都會停止回應。以下範例說明如何在 TPU 節點上執行 JAX 計算。

準備程式碼

在每個執行個體上建立名為 example.py 的檔案：

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)
EOF"

在切片上執行程式碼

在配量中的每個 TPU VM 上執行 example.py 程式：

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

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

global device count: 8
local device count: 4
pmap result: [8. 8. 8. 8.]

清除所用資源

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

或者，如要保留專案，可以只刪除 MIG 和群組中的所有 VM，請使用 gcloud compute instance-groups managed delete 指令：

gcloud compute instance-groups managed delete quickstart-tpu-mig --region=us-east5

後續步驟

• 瞭解 Compute Engine 中的 TPU 資源
• 使用其他用量方案建立多主機 TPU 配量
• 瞭解 MIG 中的 TPU 執行個體