使用 XProf 進行代管剖析

使用 ML Diagnostics 工作負載監控、ML Diagnostics CLI 或 ML Diagnostics SDK 時，可以使用代管 XProf 分析功能。代管 XProf 是 XProf 執行個體，具備可擴充的後端，由機器學習診斷平台在您的Google Cloud 專案中佈建及管理。自行託管的 XProf 則需要您手動部署及維護自己的 XProf 執行個體，才能查看設定檔。

XProf 是開放原始碼剖析和效能分析工具，適用於機器學習 (ML) 工作負載，也是 OpenXLA 生態系統的一部分。與自行託管的 XProf 相比，代管 XProf 提供下列功能：

• 無須設定：不必手動安裝或設定 XProf 視覺化伺服器或其依附元件。
• 集中存取多個執行作業的設定檔工作階段：每個專案使用單一代管 XProf 執行個體，載入不同機器學習執行作業的設定檔工作階段。
• 自動調度資源和快速載入：當多位使用者同時查看設定檔時，可加快大型設定檔的載入速度，並自動調度後端資源。
• 內建協作和安全功能：為設定檔工作階段產生共用連結，同時保護剖析執行個體免受安全漏洞影響。
• 多種設定檔擷取模式：支援多主機擷取、隨選剖析、快照模式下的持續剖析，以及從工作負載監控自動擷取設定檔。
• 無須支付額外基礎架構費用：與自行託管 XProf 執行個體不同，我們提供代管 XProf 運算資源，不收取額外費用。

部署路徑

您可以使用 ML Diagnostics SDK、工作負載監控或 ML Diagnostics CLI，部署及啟用代管 XProf。

ML Diagnostics 平台會為每個專案部署一個代管 XProf 執行個體，並搭配可擴充的後端。您可以使用這個執行個體，查看專案中所有機器學習執行作業擷取的所有設定檔工作階段。如要擷取設定檔，也需要與機器學習執行作業相關聯的 Cloud Storage bucket，才能儲存擷取的設定檔資料。

您可以透過下列任一方式啟用代管 XProf，並將 Cloud Storage bucket 與機器學習執行作業建立關聯：

• ML Diagnostics SDK：使用 SDK 部署第一個機器學習執行作業時，ML Diagnostics 平台會自動在專案中啟用及部署代管 XProf 執行個體。在 SDK 中定義機器學習執行作業時，請設定 gcs_path 參數，指定用於儲存設定檔的 Cloud Storage bucket。
• 工作負載監控： 工作負載監控功能預設不會啟用代管 XProf。相反地，當自動設定檔擷取功能首次在專案中觸發設定檔工作階段時，機器學習診斷平台會自動佈建並啟用代管 XProf 執行個體。根據預設，機器學習診斷也會在專案中佈建區域 Cloud Storage bucket。
• 機器學習診斷 CLI：您可以更新任何機器學習執行作業，啟用 XProf 工具，在專案中部署代管 XProf 執行個體，並將 Cloud Storage bucket 與機器學習執行作業建立關聯。

擷取設定檔

機器學習診斷支援四種擷取 XProf 設定檔的方法：程式輔助擷取、隨選剖析、快照模式下的持續剖析，以及從工作負載監控自動擷取設定檔。這些方法的差異在於觸發擷取的方式，以及記錄的時間範圍。

下表比較了設定檔擷取方法：

擷取方法	說明	支援的架構
程式輔助擷取	直接在工作負載程式碼中嵌入分析指令，在指定的訓練步驟或函式開始及停止記錄設定檔資料。 如要重現已知訓練步驟或特定程式碼區段的設定檔，請使用這個方法。	JAX、vLLM、Tunix、TorchTPU
隨選剖析	在工作負載執行期間，透過 Google Cloud 控制台或 gcloud CLI 觸發即時設定檔擷取，並指定擷取時間長度。 如要即時檢查執行中的工作負載，而不修改步驟範圍或重新啟動工作負載，請使用這個方法。	JAX、vLLM、Tunix、TorchTPU
快照模式下的持續剖析 (預先發布版)	在 TPU 主機的循環緩衝區中，持續記錄 2 GB 的裝置追蹤記錄設定檔，並在您觸發快照時，將最新的 2 GB 緩衝區儲存至 Cloud Storage bucket。 如要檢查導致間歇性速度變慢或異常的工作負載活動，請使用這個方法。	JAX、vLLM、Tunix
透過工作負載監控功能自動擷取設定檔 (預先發布版)	當工作負載監控效能降低分析器回報 "bad" 狀態時，自動觸發 XProf 設定檔擷取作業。您可以設定自動擷取設定檔，觸發隨選擷取，或在啟用持續剖析功能時，觸發持續剖析快照。 如果您希望機器學習診斷在效能降低時擷取追蹤記錄資料，而不必等待您手動觸發擷取，請使用這個方法。	JAX、vLLM、Tunix、TorchTPU

程式輔助擷取

透過程式輔助擷取功能，您可以直接在工作負載程式碼中嵌入分析指令，明確指定何時開始及停止記錄設定檔資料。凡是寫入機器學習執行作業 gcs_path 參數指定 Cloud Storage bucket 的設定檔，都會自動顯示為專案代管 XProf 執行個體中的設定檔工作階段，您可以在Google Cloud 控制台或使用 Google Cloud CLI 查看。

您可以透過下列任一方式設定程式輔助擷取：

• ML Diagnostics SDK：使用 SDK 中的 prof 方法。這些方法會包裝底層架構層級的分析 API，因此分析程式碼在 JAX、vLLM、Tunix 和 TorchTPU 之間會保持一致。
• 架構原生剖析 API：直接呼叫機器學習架構的內建剖析方法，並將輸出內容儲存至機器學習執行的 Cloud Storage bucket：
  • JAX：剖析運算
  • vLLM：在 Cloud TPU 上剖析 vLLM
  • TorchTPU：使用 XProf 進行效能分析

隨選剖析

注意：TorchTPU 不支援直接透過架構啟用隨選分析。如要搭配 TorchTPU 使用隨選剖析功能，請使用 ML Diagnostics SDK。

隨選剖析可讓您在工作負載執行時即時觸發設定檔擷取。觸發隨選擷取功能後，系統會在您指定的期間內記錄設定檔。您可以直接在工作負載程式碼中啟用隨選剖析功能，或使用 ML Diagnostics SDK 啟用。

必要條件

擷取隨選設定檔前，請先完成下列設定步驟：

1. 在工作負載的通訊埠 9999 上啟動 XProf 伺服器，以便回應隨選擷取觸發條件：

   • 直接在 JAX 中：在通訊埠 9999 上呼叫 profiler.start_server 方法：

     profiler.start_server(port=9999)
     

   • 機器學習診斷 SDK：定義機器學習執行作業時，請在 machinelearning_run 中設定 on_demand_xprof=True。

2. 設定叢集自動調度管理工具或運算環境：

   • Google Kubernetes Engine (GKE) 或自行管理的 Kubernetes：部署運算子和 Webhook，或 ML 診斷 DaemonSet。
   • Cluster Director 代管的 Slurm 或自行管理的 Slurm：使用 Slurm prolog 指令碼部署 ML Diagnostics daemon。
   • Compute Engine：使用執行個體開機指令碼部署機器學習診斷 daemon。

擷取隨選剖析檔

您可以透過Google Cloud 控制台或 gcloud CLI 觸發隨選設定檔擷取。觸發擷取時，請指定下列參數：

• 時間長度：擷取時間長度，以毫秒或秒為單位。可擷取的最大時間長度取決於工作負載特性，以及 2 GB 設定檔限制內可容納的資料量。
• 目標：要擷取設定檔的目標節點：
  • Google Cloud 控制台：系統會自動填入執行工作負載的所有節點，您可以選取所有或部分節點。
  • gcloud CLI：列出工作負載的分析器目標，然後觸發擷取所有或部分目標。
• 裝置、主機和 Python 追蹤器層級：要擷取的裝置、主機和 Python 追蹤記錄詳細程度。

快照模式下的持續剖析

預覽

這項功能適用《服務專屬條款》中「一般服務條款」一節的《正式發布前產品條款》。正式發布前功能是依「原樣」提供，支援服務可能受限。 詳情請參閱推出階段說明。

啟用連續剖析功能後，機器學習診斷平台會在 TPU 主機的循環緩衝區中，持續記錄 2 GB 的設定檔。觸發快照時，機器學習診斷會將緩衝區中最新的 2 GB 設定檔，儲存為 Cloud Storage bucket 中的設定檔工作階段。您可以透過這個快照，檢查問題發生前的工作負載活動。持續剖析功能僅支援裝置追蹤，不會收集主機或 Python 追蹤記錄。

注意：TorchTPU 架構不支援持續剖析。

必要條件

開始持續剖析前，請確認環境符合下列需求條件：

• 支援的程式庫版本：
  • libtpu 0.0.43 以上版本。
  • jax 0.11.0 以上版本。
• 叢集自動調度管理工具或運算環境設定：
  • GKE 或自行管理的 Kubernetes：部署 ML 診斷 DaemonSet。
  • Cluster Director 代管的 Slurm 或自行管理的 Slurm：使用 Slurm prolog 指令碼部署機器學習診斷 Daemon。
  • Compute Engine：使用執行個體開機指令碼部署機器學習診斷 daemon。

呼叫 profiler.start_server(port=9999)，在工作負載程式碼中啟動 XProf 伺服器，或在使用 ML Diagnostics SDK 時，在 machinelearning_run 中設定 on_demand_xprof=True。

開始持續剖析

您可以使用 gcloud CLI 或在 Google Cloud 控制台中，直接在工作負載程式碼中啟動持續剖析。啟動後，持續剖析功能會在指定 TPU 主機的循環緩衝區中，記錄 2 GB 的設定檔。

使用 gcloud CLI 或Google Cloud 控制台啟動持續剖析時，請指定下列參數：

• Targets：要持續剖析的目標節點：
  • Google Cloud 控制台：系統會自動填入執行工作負載的所有節點，您可以選取所有或部分節點。
  • gcloud CLI：列出工作負載的分析器目標，然後開始對所有或部分目標進行持續分析。
• 裝置追蹤器層級：持續擷取的裝置追蹤層級。持續剖析期間，主機和 Python 追蹤功能會停用。

您可以直接從 Python 工作負載程式碼、gcloud CLI 或 Google Cloud 控制台啟動連續剖析。

Python 工作負載程式碼

如要直接從 Python 工作負載程式碼啟動持續剖析，請呼叫 start_continuous_profiling 函式：

from xprof.api import continuous_profiling_snapshot
continuous_profiling_snapshot.start_continuous_profiling('localhost:9999', {})

詳情請參閱「XProf JAX 分析」。

gcloud CLI

如要從 gcloud CLI 啟動持續剖析，請列出有效的 Profiler 目標，並使用下列指令在這些節點上啟動持續剖析：

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

# 1. Discover active ProfilerTarget IDs on the ML run:
TARGETS=$(gcloud alpha mldiagnostics profiler-target list \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --format="csv[no-heading](name.basename())" | paste -sd, -)

# 2. Start continuous profiling across the target nodes:
gcloud alpha mldiagnostics machine-learning-run start-continuous-profiling \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="${TARGETS}" \
  --device-tracer-level=enabled

更改下列內容：

• PROJECT_ID：您的 Google Cloud 專案 ID。
• LOCATION：機器學習執行作業所在的區域。
• RUN_ID：機器學習執行的 ID。

Google Cloud 控制台

如要在 Google Cloud 控制台開始持續剖析，請完成下列步驟：

1. 在 Google Cloud 控制台，前往「Cluster Director」或「GKE」>「AI/機器學習」>「執行診斷」，然後選取機器學習執行作業。
2. 開啟「剖析檔」分頁。
3. 在頂端動作列中，按一下「Start continuous profiling」(開始持續剖析)。
4. 選取要剖析的節點和裝置追蹤器層級，然後按一下「Start」。

持續剖析開始後，「Capturing continuous profiles」橫幅會將工作階段的開始時間顯示為 Since START_TIME，動作列則會顯示「Take snapshot」和「Stop continuous profiling」。

擷取快照

持續剖析期間，您可以擷取快照，將循環緩衝區中最新的 2 GB 設定檔，儲存為 Cloud Storage bucket 中的設定檔工作階段。您可以使用 gcloud CLI 或 Google Cloud 控制台擷取快照。

擷取連續剖析快照時，請指定下列參數：

• 目標：執行持續剖析作業的所有或部分目標節點：
  • Google Cloud 控制台：系統會自動填入所有設為持續剖析的節點，您可以選取所有節點或部分節點。
  • gcloud CLI：指定執行持續剖析的所有或部分剖析器目標。
• Profiler 工作階段名稱：擷取快照工作階段的名稱。

您可以使用 gcloud CLI 或在Google Cloud 控制台中擷取快照。

gcloud CLI

如要使用 gcloud CLI 擷取連續剖析快照，請執行下列指令：

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

gcloud alpha mldiagnostics profiler-session capture "SNAPSHOT_NAME" \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="TARGET_NODES" \
  --kind=snapshot

更改下列內容：

• PROJECT_ID：您的 Google Cloud 專案 ID。
• LOCATION：機器學習執行作業所在的區域。
• RUN_ID：機器學習執行的 ID。
• SNAPSHOT_NAME：擷取的設定檔工作階段名稱。
• TARGET_NODES：以半形逗號分隔的目標節點清單，用於擷取。

Google Cloud 控制台

如要在 Google Cloud 控制台中擷取連續剖析快照，請完成下列步驟：

1. 在 Google Cloud 控制台中，依序前往「Cluster Director」或「GKE」>「AI/ML」>「Run Diagnostics」，選取機器學習執行作業，然後開啟「Profiles」分頁。
2. 在頂端動作列中，按一下「建立快照」。
3. 在「Capture continuous profiling snapshot」(擷取連續剖析快照) 窗格中，輸入設定檔名稱，選取目標節點，然後按一下「Submit」(提交)。

快照會顯示在「Profiles」表格中。如果連結未立即顯示，請按一下「Refresh」(重新整理)。

透過工作負載監控功能自動擷取設定檔

預覽

這項功能適用《服務專屬條款》中「一般服務條款」一節的《正式發布前產品條款》。正式發布前功能是依「原樣」提供，支援服務可能受限。 詳情請參閱推出階段說明。

透過工作負載監控功能自動擷取設定檔，當效能降低分析器回報 "bad" 狀態時，機器學習診斷會自動擷取 XProf 設定檔，不必等到問題發生後，才手動觸發隨選擷取。您可以設定自動設定檔擷取功能，觸發隨選擷取，或在啟用持續剖析功能時，觸發持續剖析快照。這個程序可讓您使用事件發生時擷取的設定檔，診斷問題的根本原因。

如需設定操作說明和必要條件，請參閱工作負載監控說明文件中的「自動擷取設定檔」。

代理式 AI 設定檔分析

預覽

這項功能適用《服務專屬條款》中「一般服務條款」一節的《正式發布前產品條款》。正式發布前功能是依「原樣」提供，支援服務可能受限。 詳情請參閱推出階段說明。

代理式 AI 設定檔分析功能會自動分析為 TPU 工作負載擷取的 XProf 效能設定檔，並根據原始 XProf 追蹤記錄資料產生 HTML 分析報表。分析代理會使用 Vertex AI 上的 Gemini 模型，執行 XProf 剖析工具、總結效能瓶頸，並建議程式碼或設定變更。啟用必要的 API 後，機器學習診斷功能會為平台擷取的每個設定檔，生成代理式 AI 設定檔分析報表。

代理式 AI 設定檔分析使用的基礎工具和功能，與開放原始碼 XProf 技能存放區相同。ML Diagnostics 不會直接使用技能檔案，或安裝開放原始碼技能所需的 XProf CLI，而是匯入基礎 XProf 工具實作，讓代理可以直接呼叫。因此，機器學習診斷中的分析輸出內容，等同於在本機 xplane 檔案上，使用 XProf CLI 執行開放原始碼 XProf 技能。

必要條件

使用代理式 AI 設定檔分析功能前，請先在 Google Cloud 專案中完成下列設定步驟。

啟用必要的 API

啟用 Vertex AI API (aiplatform.googleapis.com)，讓分析代理在專案中執行 Gemini 模型：

• Google Cloud 控制台: 在API 程式庫中，啟用Vertex AI API。

• gcloud CLI：執行下列指令：

  gcloud services enable aiplatform.googleapis.com
  

授予 IAM 角色

授予使用者帳戶下列其中一個 Identity and Access Management (IAM) 角色：

• Cluster Director 檢視者 (roles/hypercomputecluster.viewer)：可唯讀存取 Cluster Director 和機器學習診斷資源，包括 Google Cloud 控制台「Run Diagnostics」>「Profiles」下方產生的分析報表。
• 自訂角色 (hypercomputecluster.locations.analysis)：新增至自訂角色時，授予權限，可在 Google Cloud 控制台的「Run Diagnostics」>「Profiles」下方，查看產生的分析報表。

運作方式

當程式輔助、隨選、持續或自動設定檔擷取完成時，獨立的背景工作會在 XProf 設定檔上執行下列分析工作流程：

1. 廣泛收集資料：同時執行 15 項 XProf 工具，收集健康狀態快照，包括步驟時間、硬體使用率、記憶體用量、roofline 限制、最慢作業和編譯器效率不彰。這個步驟會將每個訓練步驟的時間正規化，並區分正常作業和瓶頸。
2. Focused AI investigation：查看收集到的健康狀態匯報。如果工作負載的運算使用率高、閒置和通訊負擔低，且記憶體不足風險低，代理就會略過其他工具呼叫，並產生 OPTIMAL 報告。如果偵測到瓶頸，代理會提出假設，並使用其他 XProf 工具執行最多四項指定後續檢查，檢查運算圖中的周圍作業、主機之間的時間軸延遲，以及自訂核心。
3. 報表生成：結合 XProf 的指標與 AI 代理的根本原因分析和建議修正。產生的summary.html報表會儲存至 Cloud Storage bucket，並顯示在Google Cloud 控制台的「Run Diagnostics」>「Profiles」下方。

這項工作流程可確保設定檔資料和分析結果在專案中保持私密。原始分析器追蹤記錄和產生的 HTML 報表會儲存在 Cloud Storage bucket，所有 AI 模型呼叫都會透過 Google Cloud 專案中的 Vertex AI 執行。

XProf 工具

下表說明代理程式在分析期間執行的 XProf 工具：

類別	XProf 工具	說明
工作負載總覽和步驟時間	get_xprof_overview、get_step_trace、get_kpi_metrics、get_device_information、get_hosts、get_utilization_viewer、get_smart_suggestions	這些工具會在廣泛資料收集期間執行，而 get_hosts 工具也會在重點 AI 調查期間執行，以檢查下列項目： 硬體類型、主機和晶片數量，以及硬體上限 平均步驟時間，以及運算、網路通訊和主機或資料載入閒置時間的細目 矩陣乘法單元 (MXU)、TensorCore 和記憶體頻寬使用率，以及內建的 XProf 建議
Roofline 和最慢作業	get_roofline_model，get_hlo_op_profile	這些工具會在廣泛資料收集和重點 AI 調查期間執行，檢查下列項目： 工作負載或特定作業是否受限於運算或記憶體頻寬 依運算類型 (例如矩陣乘法、注意力機制和集體運算) 劃分的時間，以及 Python 原始碼行的連結
編譯器效率不彰偵測工具	detect_layout_mismatch_copies、detect_unfused_reshapes、detect_unnecessary_convert_reduce、list_hlo_modules	這些工具會在廣泛收集資料期間執行，檢查下列項目： 運算之間的張量版面配置不相符，導致額外的資料副本 未融合的 reshape、transpose 或 copy 作業，會浪費記憶體頻寬 縮減作業中不必要的精確度轉換，例如 bf16 轉換為 f32
記憶體和 OOM 風險	get_memory_profile，get_peak_allocations	這些工具會在廣泛收集資料期間執行，檢查下列項目： 裝置上的高頻寬記憶體 (HBM) 尖峰用量 (與總容量和記憶體片段相比) 記憶體不足 (OOM) 風險是否為LOW、MEDIUM或HIGH，以及記憶體中最大的張量
運算圖探索工具	get_hlo_neighborhood	這項工具會在 AI 專注調查期間執行，檢查 XLA 圖中慢速作業前後的作業，並找出阻礙編譯器融合或強制移動資料的原因。
時間軸和核心時間	aggregate_xplane_events，get_kernel_stats	這些工具會在 AI 專注調查期間執行，檢查下列項目： 跨主機的時間軸統計資料，可區分資料載入器停滯和網路通訊延遲 個別裝置核心的執行時間和呼叫次數
自訂核心分析器	llo_load，llo_query	這些工具會在 AI 專項調查期間執行，檢查自訂核心 (例如 Pallas 或 Mosaic) 的低階執行追蹤記錄，這些核心在標準高階剖析中會顯示為不透明。

AI 目標對象資料分析報表

擷取的設定檔分析完成後，機器學習診斷會將產生的summary.html報表儲存至 Cloud Storage bucket，並在 Google Cloud 控制台的「Run Diagnostics」>「Profiles」下方顯示。您可以透過這份報告查看工作負載的整體健康狀態、檢查主要效能瓶頸，以及找出建議的程式碼或設定變更。

每份產生的報表都包含下列部分：

• 分析摘要：顯示 OPTIMAL 或 NEEDS_ATTENTION 的整體狀態、3 到 5 句的摘要，以及步驟時間、記憶體用量和 Roofline 效率等關鍵成效指標。
• 優先處理瓶頸和可行的修正措施：列出發現項目，並依優先順序排列，從 P0 (重大問題) 到 P2 (次要問題)。每項發現項目都包含：
  • 發生什麼事：說明根本原因，包括張量形狀和版面配置 (如適用)。
  • 證據：直接引用 XProf 指標和作業，指出瓶頸的測量位置。
  • 預期影響和修正：受影響的步驟時間或記憶體量，以及建議的架構、資料分割或編譯器變更。
• 詳細的 XProf 工具分析：提供下列分析的支援資料和詳細細目：
  • Roofline 分析
  • 記憶體分析
  • 高階最佳化工具 (HLO) 作業設定檔分析
  • 核心執行統計資料
  • 自訂核心和低階最佳化工具分析