使用 XProf 剖析 JAX 運算

XProf 是取得及顯示程式效能追蹤記錄和設定檔的絕佳方式，包括 GPU 和 TPU 上的活動。最終結果如下所示：

XProf 範例

以程式輔助方式擷取

您可以檢測程式碼，透過 jax.profiler.start_trace 和 jax.profiler.stop_trace 方法擷取 JAX 程式碼的剖析器追蹤記錄。呼叫 jax.profiler.start_trace，並提供要寫入追蹤記錄檔案的目錄。這個目錄應與用於啟動 XProf 的目錄相同。--logdir接著，您可以使用 XProf 查看追蹤記錄。

舉例來說，如要擷取分析器追蹤記錄：

import jax

jax.profiler.start_trace("/tmp/profile-data")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

請注意 jax.block_until_ready 呼叫。我們會使用這項功能，確保追蹤記錄會擷取裝置端執行作業。詳情請參閱「非同步調度」，瞭解為何需要這麼做。

您也可以使用 jax.profiler.trace 內容管理工具，做為 start_trace 和 stop_trace 的替代方案：

import jax

with jax.profiler.trace("/tmp/profile-data"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

查看追蹤記錄

擷取追蹤記錄後，您可以使用 XProf UI 查看。

您可以指向記錄目錄，直接使用獨立的 XProf 指令啟動分析器 UI：

$ xprof --port=8791 /tmp/profile-data
Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
  Worker Service Address: 0.0.0.0:50051
  Hide Capture Button: False
XProf at http://localhost:8791/ (Press CTRL+C to quit)

前往提供的網址 (例如 http://localhost:8791/) 即可在瀏覽器中查看商家檔案。

左側的「Sessions」下拉式選單會顯示可用的追蹤記錄。選取感興趣的階段作業，然後在「工具」下拉式選單下方選取「追蹤檢視器」。現在應該會看到執行時間軸。您可以使用 WASD 鍵瀏覽追蹤記錄，然後按一下或拖曳選取事件，查看更多詳細資料。如要進一步瞭解如何使用追蹤記錄檢視器，請參閱追蹤記錄檢視器工具說明文件。

透過 XProf 手動擷取

以下說明如何從執行中的程式擷取手動觸發的 N 秒追蹤記錄。

啟動 XProf 伺服器：
```
xprof --logdir /tmp/profile-data/
```
您應該可以在 <http://localhost:8791/> 載入 XProf。您可以使用 --port 旗標指定其他通訊埠。
在要剖析的 Python 程式或程序中，於開頭附近新增下列內容：
```
import jax.profiler
jax.profiler.start_server(9999)
```
這會啟動 XProf 連線的剖析器伺服器。您必須先執行剖析器伺服器，才能進行下一個步驟。使用完伺服器後，可以呼叫 jax.profiler.stop_server() 將其關機。

如要剖析長時間執行的程式片段 (例如長時間的訓練迴圈)，您可以將這項設定放在程式開頭，然後照常啟動程式。如要分析短程式 (例如微基準)，其中一個方法是在 IPython 殼層中啟動分析器伺服器，並在下一個步驟中啟動擷取作業後，使用 %run 執行短程式。您也可以在程式開頭啟動剖析器伺服器，並使用 time.sleep()，爭取足夠的時間啟動擷取作業。
開啟 <http://localhost:8791/>，然後按一下左上方的「擷取設定檔」按鈕。輸入「localhost:9999」做為設定檔服務網址 (這是您在上一個步驟中啟動的剖析器伺服器位址)。輸入要分析的毫秒數，然後按一下「CAPTURE」。
如果想剖析的程式碼尚未執行 (例如在 Python 殼層中啟動剖析器伺服器)，請在擷取作業執行期間執行該程式碼。
擷取完成後，XProf 應會自動重新整理。(並非所有 XProf 分析功能都與 JAX 連結，因此一開始可能看起來沒有擷取任何內容)。在左側的「工具」下方，選取「追蹤檢視器」。

現在應該會看到執行時間軸。您可以使用 WASD 鍵瀏覽追蹤記錄，然後按一下或拖曳選取事件，即可在底部查看更多詳細資料。如要進一步瞭解如何使用追蹤記錄檢視器，請參閱追蹤記錄檢視器工具說明文件。

XProf 和 TensorBoard

XProf 是 TensorBoard 中剖析和追蹤記錄擷取功能所用的基礎工具。只要安裝 xprof，TensorBoard 就會顯示「Profile」分頁。只要啟動時指向相同的記錄目錄，使用這個選項就等同於獨立啟動 XProf。包括擷取、分析及檢視設定檔的功能。XProf 取代了先前建議的 tensorboard_plugin_profile 功能。

$ tensorboard --logdir=/tmp/profile-data
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.19.0 at http://localhost:6006/ (Press CTRL+C to quit)

新增自訂追蹤事件

根據預設，追蹤記錄檢視器中的事件大多是低階內部 JAX 函式。您可以在程式碼中使用 jax.profiler.TraceAnnotation 和 jax.profiler.annotate_function 新增自己的事件和函式。

設定分析器選項

start_trace 方法接受選用的 profiler_options 參數，可精細控管剖析器的行為。這個參數應為 jax.profiler.ProfileOptions 的例項。

舉例來說，如要停用所有 Python 和主機追蹤記錄，請執行下列指令：

import jax

options = jax.profiler.ProfileOptions()
options.python_tracer_level = 0
options.host_tracer_level = 0
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

一般選項

host_tracer_level：設定主機端活動的追蹤層級。

支援的值：
- 0：完全停用主機 (CPU) 追蹤。
- 1：僅追蹤使用者插碼的 TraceMe 事件。
- 2：包含第 1 級追蹤記錄，以及高階程式執行詳細資料，例如耗用大量資源的 XLA 作業 (預設)。
- 3：包括第 2 級追蹤記錄，以及更詳細的低階程式執行詳細資料，例如便宜的 XLA 作業。
device_tracer_level：控制是否啟用裝置追蹤功能。

支援的值：
- 0：停用裝置追蹤功能。
- 1：啟用裝置追蹤功能 (預設)。
python_tracer_level：控制是否啟用 Python 追蹤功能。

支援的值：
- 0：停用 Python 函式呼叫追蹤功能 (預設)。
- 1：啟用 Python 追蹤。

進階設定選項

TPU 選項

tpu_trace_mode：指定 TPU 追蹤模式。

支援的值：
- TRACE_ONLY_HOST：這表示系統只會追蹤主機端 (CPU) 活動，不會收集裝置 (TPU/GPU) 追蹤記錄。
- TRACE_ONLY_XLA：這表示系統只會追蹤裝置上的 XLA 層級作業。
- TRACE_COMPUTE：追蹤裝置上的運算作業。
- TRACE_COMPUTE_AND_SYNC：這會追蹤裝置上的運算作業和同步事件。
如未提供「tpu_trace_mode」，trace_mode 預設為 TRACE_ONLY_XLA。
tpu_num_sparse_cores_to_trace：指定要在 TPU 上追蹤的稀疏核心數量。
tpu_num_sparse_core_tiles_to_trace：指定要在 TPU 上追蹤的每個稀疏核心內圖塊數量。
tpu_num_chips_to_profile_per_task：指定要為每個工作剖析的 TPU 晶片數量。

GPU 選項

GPU 剖析可用的選項如下：

gpu_max_callback_api_events：設定 CUPTI 回呼 API 收集的事件數量上限。預設為 2*1024*1024。
gpu_max_activity_api_events：設定 CUPTI 活動 API 收集的事件數量上限。預設為 2*1024*1024。
gpu_max_annotation_strings：設定可收集的註解字串數量上限。預設為 1024*1024。
gpu_enable_nvtx_tracking：在 CUPTI 中啟用 NVTX 追蹤功能。預設值為 False。
gpu_enable_cupti_activity_graph_trace：啟用 CUDA 圖形的 CUPTI 活動圖形追蹤功能。預設為 False。
gpu_pm_sample_counters：以半形逗號分隔的 GPU 效能監控指標字串，用於透過 CUPTI 的 PM 抽樣功能收集資料 (例如 "sm__cycles_active.avg.pct_of_peak_sustained_elapsed")。PM 抽樣功能預設為停用。如需可用指標，請參閱 NVIDIA 的 CUPTI 說明文件。
gpu_pm_sample_interval_us：設定 CUPTI PM 抽樣的抽樣間隔 (以微秒為單位)。預設為 500。
gpu_pm_sample_buffer_size_per_gpu_mb：設定每個裝置的系統記憶體緩衝區大小 (以 MB 為單位)，用於 CUPTI PM 採樣。預設值為 64 MB。支援的最大值為 4 GB。
gpu_num_chips_to_profile_per_task：指定每個工作要剖析的 GPU 裝置數量。如未指定、設為 0 或設為無效值，系統會分析所有可用的 GPU。這項功能可用於縮減追蹤記錄收集大小。
gpu_dump_graph_node_mapping：啟用後，會將 CUDA 圖形節點對應資訊傾印至追蹤記錄。預設為 False。

例如：

options = ProfileOptions()
options.advanced_configuration = {"tpu_trace_mode" : "TRACE_ONLY_HOST", "tpu_num_sparse_cores_to_trace" : 2}

如果發現任何無法辨識的鍵或選項值，則傳回 InvalidArgumentError。

使用 XProf 剖析 JAX 運算 透過集合功能整理內容 你可以依據偏好儲存及分類內容。