在 Cloud TPU VM 上剖析模型
剖析是最佳化 Cloud TPU 模型效能的主要工具之一。核心剖析工具稱為 XProf，可從 OpenXLA/XProf GitHub 存放區取得。XProf 支援所有 XLA 架構的剖析，包括 JAX、PyTorch 和 Tensorflow/Keras。
擷取設定檔
如要使用 XProf，請先擷取模型設定檔。擷取設定檔的方法有兩種：
- 程式輔助擷取
- 隨選擷取 (又稱手動擷取)
使用程式輔助擷取時，您必須為模型程式碼加上註解，指定要擷取設定檔的位置。一般來說，您會擷取幾個訓練步驟的設定檔，或剖析模型中特定程式碼區塊。在 JAX、PyTorch 和 TensorFlow 等不同架構中，擷取追蹤記錄的方式也不同，包括透過 API 啟動或停止追蹤，或使用內容管理工具。在 MaxText 等高階架構中，只要在執行 MaxText 時啟用旗標 profiler=xplane，即可啟用設定檔擷取功能。
如要以臨時方式擷取設定檔，或未啟用程式輔助設定檔擷取功能，可以使用隨選設定檔擷取功能。如果執行期間模型指標出現問題，而您想擷取該時間點的設定檔，並持續一段時間，以便診斷問題，這個方法就很有幫助。
如要啟用隨選設定檔擷取功能，您仍須在程式碼中啟動 XProf 伺服器。啟用 profiler.start_server 後，系統會在機器學習工作負載上啟動 XProf 伺服器，監聽隨選擷取觸發條件，以便開始擷取設定檔。您可以透過 Tensorboard UI 或 CLI 使用 XProfiler 工具，觸發隨選設定檔擷取。
如要瞭解如何為不同框架啟用程式輔助和隨選擷取功能，請參閱：
無論是程式輔助或隨選設定檔擷取，您都必須指定擷取設定檔的儲存位置。您可以選擇將設定檔儲存在 TPU VM 的目錄中，路徑類似 /profiles/run-name，或使用 Cloud Storage，路徑類似 gs://bucket-name/run-name/。
使用上述路徑，系統會在以下位置擷取設定檔：
/profiles/run-name/plugins/profile/session1/
或
gs://bucket-name/run-name/plugins/profile/session1/。
工作階段代表模型執行期間擷取的一組設定檔。舉例來說，如果您在訓練執行作業中擷取步驟 1-3 的設定檔，並在稍後擷取步驟 8-10 的設定檔，則這些設定檔屬於同一個執行作業，但步驟 1-3 的第一個擷取作業會是 session1，而步驟 8-10 的第二個擷取作業會是 session2。
每次執行下方會以不同日期戳記標示不同工作階段。 您可以透過程式輔助或隨選方式，或結合兩者，在不同工作階段中擷取設定檔。如要比較模型訓練不同部分的設定檔，這項功能就非常實用。例如，比較訓練開始時和訓練結束時的設定檔。
使用 TensorBoard 以圖表呈現設定檔
如要在 Tensorboard 中查看 XProf 設定檔，您必須安裝 tensorboard-plugin-profile 外掛程式。
如要在 Google Cloud上以圖表呈現擷取的設定檔，建議您在 TPU VM 上使用 TensorBoard，或使用 XProfiler 工具在 VM 或 GKE Pod 上託管 TensorBoard。
如果您在 TPU VM 上擷取設定檔，可以使用該 TPU VM 上的 TensorBoard 查看擷取的設定檔。
如要在 TPU VM 上安裝 XProf 和 TensorBoard，請執行下列指令：
pip install tensorboard_plugin_profile tensorboard pip install xprof
如要查看 TensorBoard 擷取的設定檔，請執行下列指令：
$ tensorboard --logdir=profiles/run-name
或
$ tensorboard --logdir=profiles
如要從多個執行作業載入多個設定檔，請將 TensorBoard 指向包含所有執行作業和設定檔的根目錄，而非包含特定執行作業設定檔的子目錄。
Cloud-Diagnostics-XProf 程式庫，用於在 Google Cloud上以圖表呈現 TensorBoard 設定檔
cloud-diagnostics-xprof (又稱 XProfiler 程式庫) 可讓您更輕鬆地在 Google Cloud上託管 TensorBoard 並以視覺化方式呈現設定檔。詳情請參閱 cloud-diagnostics-xprof GitHub 存放區。
與在本機執行 XProf 和 TensorBoard 相比，cloud-diagnostics-xprof 存放區提供下列改善項目：
- 設定及封裝 XProf 和 TensorBoard 依附元件
- 將設定檔儲存在 Cloud Storage，以便長期保留及執行後分析 (研究人員完成執行後，系統會刪除擷取的本機設定檔)
- 在 Compute Engine VM 或 GKE Pod 上佈建 TensorBoard，快速載入大型設定檔和多個設定檔，並可根據使用者對載入速度和成本的需求變更機型
- 建立連結，與團隊成員和 Google 工程師分享設定檔並協作
- 更輕鬆地在 GKE 和 Compute Engine 上隨選剖析工作負載，選擇任何執行工作負載的主機來擷取設定檔
使用 XProfiler 程式庫前，您必須透過程式輔助方式或啟動設定檔伺服器，擷取工作負載程式碼的設定檔，以便日後隨選擷取設定檔。
如要設定 XProfiler 程式庫，您必須先設定 Google Cloud CLI 和 Python 虛擬環境，然後執行下列指令：
pip install cloud-diagnostics-xprof
這會安裝所有必要的 XProf 和 TensorBoard 依附元件。
接著，執行下列指令，建立 VM 或 GKE Pod 來託管 TensorBoard：
xprofiler create -z $ZONE -l $GCS_PATH
或
xprofiler create --GKE -z $ZONE -l $GCS_PATH
將 $ZONE 替換為任何可用區，並將 $GCS_PATH 替換為設定檔追蹤記錄的路徑。您可以指定根目錄，其中包含多次執行的所有設定檔追蹤記錄，或一次執行的特定設定檔追蹤記錄。
舉例來說，如果您在下列目錄中擷取設定檔：
gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>
您可以將 Cloud Storage 路徑設為根目錄 (gs://bucket-name)：
xprofiler create -z $ZONE -l gs://bucket-name/
在 TensorBoard UI 中，您會看到所有執行作業和工作階段的設定檔，例如 run1/session1、run1/session2、run2/session1。
根據預設，xprofiler create 會建立 Compute Engine VM，具體來說是 c4-highmem-8 VM。您可以使用 -m 旗標變更機器類型。如要在 GKE Pod (而非 Compute Engine VM) 上建立 TensorBoard 執行個體，可以將 –GKE 旗標傳遞至 xprofiler create。在 GKE Pod 上託管 TensorBoard 執行個體，可讓您更輕鬆地管理 TensorBoard，以及部署在 GKE 上的其他工作負載。
相較於本機託管的 TensorBoard，託管 TensorBoard 的 VM 或 GKE Pod 在 Google Cloud 上載入大型設定檔和多個設定檔的速度更快。根據我們的基準測試，使用預設的 c4-highmem-8 VM 首次載入時，您會在幾分鐘內看到 1 GB 載入量的設定檔。此外，您也可以根據效能和成本需求，選擇合適的機器類型。
執行 xprofiler create 後，您會看到類似下方的輸出內容：
Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.
第一個選項是連結，點選後即可在 TensorBoard 上查看 XProf 設定檔。這個連結可讓您與團隊和 Google 工程師分享設定檔，協助 Google Cloud進行效能最佳化。
您可以根據儲存設定檔資料的 Cloud Storage bucket 權限，控管哪些人可以存取連結。如果連結因故無法運作，我們也提供使用 SSH 連線至 TensorBoard 執行個體的方法，以便使用 xprofiler connect 指令查看設定檔。
如果您在工作負載程式碼中啟用了 Cloud Profiler 伺服器 (程式輔助和隨選擷取都需要)，並想執行隨選分析，可以透過下列兩種方式進行：
a. 按一下 TensorBoard 上的「Capture profile」按鈕。現在，您可以選擇任何執行工作負載的裝置主機進行分析。我們支援對在 GKE 或 Compute Engine 上執行的工作負載進行隨選擷取。
b. 使用 XProfiler 擷取指令，提供可用區、Cloud Storage bucket、架構、主機 VM/Pod 名稱，以及擷取時間長度 (以毫秒為單位) 等資訊。這項資訊與您在 TensorBoard UI 中輸入的資訊相同。
如要進一步瞭解如何使用 cloud-diagnostics-xprof 程式庫，請參閱這個 GitHub 頁面。