本頁面由 Cloud Translation API 翻譯而成。

剖析 PyTorch XLA 工作負載

剖析是一種分析及改善模型效能的方法。雖然還有許多其他因素，但有時將剖析視為計時作業，以及在裝置 (TPU) 和主機 (CPU) 上執行的程式碼部分，有助於您瞭解剖析結果。本指南將簡要說明如何剖析程式碼，以利訓練或推論。如要進一步瞭解如何分析產生的設定檔，請參閱下列指南。

建立 TPU

匯出環境變數：
```
$ export TPU_NAME=your_tpu_name
$ export ZONE=us-central2-b
$ export PROJECT_ID=project-id
$ export ACCELERATOR_TYPE=v4-8
$ export RUNTIME_VERSION=tpu-vm-v4-pt-2.0
```
匯出變數說明

TPU name

要用於 Cloud TPU 的名稱。

zone

您要建立 Cloud TPU 的區域。

project ID

您用於訓練及剖析模型的專案 ID。

accelerator-type

加速器類型會指定您要建立的 Cloud TPU 版本和大小。如要進一步瞭解各 TPU 版本支援的加速器類型，請參閱「TPU 版本」。

version

Cloud TPU 執行階段版本。匯出的變數會顯示預設值，但您也可以使用支援的設定清單中的值。

啟動 TPU 資源

$ gcloud compute tpus tpu-vm create ${TPU_NAME} \
    --zone us-central2-b \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --version ${RUNTIME_VERSION} \
    --project ${PROJECT_ID}

請使用下列指令，在 TPU 分割區的所有 TPU VM 上安裝 torch_xla。您也需要安裝訓練指令碼所需的其他依附元件。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --worker=all \
    --command="pip install torch==2.6.0 torch_xla[tpu]==2.6.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html"

使用 gcloud scp 指令，將程式碼移至 TPU VM 的主目錄。例如：

$ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}:directory/target-file --zone ${ZONE}

設定檔

您可以透過 capture_profile.py 手動擷取設定檔，也可以使用 torch_xla.debug.profiler API 在訓練指令碼中以程式輔助方式擷取設定檔。

啟動設定檔伺服器

如要擷取設定檔，訓練指令碼中必須執行設定檔伺服器。使用您選擇的通訊埠號碼啟動伺服器，例如下列指令所示的 9012。

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)

您可以在 main 函式的開頭啟動伺服器。

您現在可以按照下文所述擷取設定檔。這個指令碼會分析 TPU 裝置上發生的所有事件。

新增追蹤記錄

如果您也想在主機上剖析作業，可以在程式碼中新增 xp.StepTrace 或 xp.Trace。這些函式會追蹤主機上的 Python 程式碼。(您可以將這項測試視為在將「圖表」傳送至 TPU 裝置之前，在主機 (CPU) 上執行 Python 程式碼所需的時間。因此，這項指標最適合用於分析追蹤額外負荷。您可以在程式碼處理資料批次的訓練迴圈中加入此值，例如：

for step, batch in enumerate(train_dataloader):
   with xp.StepTrace('Training_step', step_num=step):
   ...

或將程式碼的個別部分包裝在

with xp.Trace('loss'):
   loss = ...

如果您使用的是 Lighting，可以略過新增追蹤記錄的步驟，因為程式碼的某些部分會自動執行這項操作。不過，如果您想新增其他追蹤記錄，歡迎在訓練迴圈中插入這些記錄。

您可以在初始編譯後擷取裝置活動，等待模型開始訓練或推論步驟。

手動擷取

Pytorch XLA 存放區的 capture_profile.py 指令碼可快速擷取設定檔。方法是將擷取設定檔檔案直接複製到 TPU VM。下列指令會將檔案複製到主目錄。

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone us-central2-b \
    --worker=all \
    --command="wget https://raw.githubusercontent.com/pytorch/xla/master/scripts/capture_profile.py"

在訓練執行期間，執行下列指令來擷取設定檔：

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone us-central2-b \
    --worker=all \
    --command="python3 capture_profile.py --service_addr "localhost:9012" --logdir ~/profiles/ --duration_ms 2000"

這個指令會將 .xplane.pb 檔案儲存在 logdir 中。您可以將記錄目錄 ~/profiles/ 變更為偏好的位置和名稱。您也可以直接將檔案儲存至 Cloud Storage 值區。如要這麼做，請將 logdir 設為 gs://your_bucket_name/。

程式輔助擷取

您可以使用訓練指令碼中的 torch_xla.debug.profiler.trace_detached API，讓訓練指令碼自動觸發設定檔，而非透過觸發指令碼手動擷取設定檔。

舉例來說，如要自動擷取特定 epoch 和步驟的設定檔，您可以設定訓練指令碼，以便使用 PROFILE_STEP、PROFILE_EPOCH 和 PROFILE_LOGDIR 環境變數：

import os
import torch_xla.debug.profiler as xp

# Within the training script, read the step and epoch to profile from the
# environment.
profile_step = int(os.environ.get('PROFILE_STEP', -1))
profile_epoch = int(os.environ.get('PROFILE_EPOCH', -1))
...

for epoch in range(num_epoch):
   ...
   for step, data in enumerate(epoch_dataloader):
      if epoch == profile_epoch and step == profile_step:
         profile_logdir = os.environ['PROFILE_LOGDIR']
         # Use trace_detached to capture the profile from a background thread
         xp.trace_detached('localhost:9012', profile_logdir)
      ...

這會將 .xplane.pb 檔案儲存在 PROFILE_LOGDIR 環境變數指定的目錄中。

TensorBoard 中的分析

如要進一步分析設定檔，您可以在同一部或其他機器上使用 TensorBoard 搭配 TPU TensorBoard 外掛程式 (建議做法)。

如要在遠端電腦上執行 TensorBoard，請使用 SSH 連線至該電腦，並啟用通訊埠轉送功能。例如，假設使用者要求系統將文字從英文翻譯成法文

$ ssh -L 6006:localhost:6006 remote server address

或

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --ssh-flag="-4 -L 6006:localhost:6006"

在遠端電腦上安裝必要套件，並啟動 TensorBoard (假設您在該電腦上有 ~/profiles/ 下的設定檔)。如果您將設定檔儲存在其他目錄或 Cloud Storage 值區中，請務必正確指定路徑，例如 gs://your_bucket_name/profiles。

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile

(vm)$ tensorboard --logdir ~/profiles/ --port 6006

(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

執行 TensorBoard

在本機瀏覽器中前往 http://localhost:6006/，然後從下拉式選單中選擇 PROFILE 來載入設定檔。

如要瞭解 TensorBoard 工具和如何解讀輸出內容，請參閱「在 Cloud TPU VM 上剖析模型」。

TensorBoard 剖析資料擷取頁面