使用 TensorFlow Profiler 優化 TensorFlow GPU 效能

總覽

本指南將說明如何搭配 TensorBoard 使用 TensorFlow Profiler，以深入瞭解 GPU 並發揮 GPU 的最大效能，以及在一個或多個 GPU 未充分利用時進行偵錯。

如果您是 Profiler 新手

從「TensorFlow Profiler：分析模型效能」筆記本開始，其中包含 Keras 範例和 TensorBoard。
透過「使用 Profiler 優化 TensorFlow 效能」指南，瞭解各種可用於在主機 (CPU) 上優化 TensorFlow 效能的分析工具和方法。

請記住，將運算卸載到 GPU 可能並非總是帶來好處，對於小型模型尤其如此。可能會因為以下原因而產生額外負擔：

主機 (CPU) 和裝置 (GPU) 之間的資料傳輸；以及
主機啟動 GPU 核心時涉及的延遲。

效能最佳化工作流程

本指南概述如何偵錯效能問題，從單一 GPU 開始，然後移至具有多個 GPU 的單一主機。

建議依以下順序偵錯效能問題：

優化和偵錯單一 GPU 的效能
1. 檢查輸入管線是否為瓶頸。
2. 偵錯單一 GPU 的效能。
3. 啟用混合精度（使用 fp16 (float16)），並視需要啟用 XLA。
優化和偵錯多 GPU 單一主機的效能。

例如，如果您使用 TensorFlow 分散式策略在具有多個 GPU 的單一主機上訓練模型，並注意到 GPU 使用率欠佳，您應先優化和偵錯單一 GPU 的效能，再偵錯多 GPU 系統。

作為在 GPU 上獲得高效能程式碼的基準，本指南假設您已在使用 tf.function。Keras Model.compile 和 Model.fit API 會在底層自動使用 tf.function。當使用 tf.GradientTape 撰寫自訂訓練迴圈時，請參閱「使用 tf.function 獲得更佳效能」以瞭解如何啟用 tf.function。

以下各節討論針對上述各種情境建議的方法，以協助找出和修正效能瓶頸。

1. 優化單一 GPU 的效能

在理想情況下，您的程式應具有高 GPU 使用率、最少的主機 (CPU) 到裝置 (GPU) 通訊，且輸入管線沒有額外負擔。

分析效能的第一步是取得使用單一 GPU 執行的模型的設定檔。

TensorBoard 的 Profiler「總覽頁面」會顯示模型在設定檔執行期間的頂層檢視畫面，可讓您瞭解您的程式與理想情境的差距。

TensorFlow Profiler Overview Page

需要注意總覽頁面的主要數字包括：

實際裝置執行佔步驟時間的比例
放置在裝置與主機上的運算百分比
有多少核心使用 fp16

達到最佳效能表示在所有三種情況下都應最大化這些數字。若要深入瞭解您的程式，您需要熟悉 TensorBoard 的 Profiler「追蹤檢視器」。以下各節說明診斷效能瓶頸時應注意的一些常見追蹤檢視器模式。

以下是使用單一 GPU 執行的模型追蹤檢視畫面圖片。從「TensorFlow 名稱範圍」和「TensorFlow 運算」區段中，您可以識別模型的不同部分，例如前向傳遞、損失函式、反向傳遞/梯度計算，以及最佳化工具權重更新。您也可以讓在 GPU 上執行的運算位於每個「串流」旁邊，這些串流是指 CUDA 串流。每個串流都用於特定工作。在此追蹤中，「串流 #118」用於啟動運算核心和裝置對裝置複製。「串流 #119」用於主機對裝置複製，「串流 #120」用於裝置對主機複製。

以下追蹤顯示高效能模型的常見特徵。

例如，GPU 運算時間軸 (「串流 #118」) 看起來「繁忙」，間隙非常少。主機到裝置 (「串流 #119」) 和裝置到主機 (「串流 #120」) 的複製最少，且步驟之間的間隙也很小。當您為程式執行 Profiler 時，您可能無法在追蹤檢視畫面中識別出這些理想的特徵。本指南的其餘部分涵蓋常見情境以及如何修正這些情境。

1. 偵錯輸入管線

GPU 效能偵錯的第一步是判斷您的程式是否受輸入限制。找出此問題的最簡單方法是使用 TensorBoard 上的 Profiler「輸入管線分析器」，該分析器會提供輸入管線中花費時間的總覽。

如果您的輸入管線對步驟時間有顯著影響，您可以採取以下潛在措施：

您可以使用 tf.data 專用的指南，以瞭解如何偵錯輸入管線。
檢查輸入管線是否為瓶頸的另一個快速方法是使用不需要任何預先處理的隨機產生輸入資料。以下範例說明將此技術用於 ResNet 模型。如果輸入管線是最佳的，您應該在使用真實資料和使用產生的隨機/合成資料時體驗到類似的效能。合成資料案例中唯一的額外負擔將是由於輸入資料複製，而輸入資料複製可以再次預先擷取和最佳化。

此外，請參閱最佳化輸入資料管線的最佳做法。

2. 偵錯單一 GPU 的效能

有多個因素可能導致 GPU 使用率偏低。以下是在查看追蹤檢視器時常見的一些情境和潛在解決方案。

1. 分析步驟之間的間隙

當您的程式未以最佳方式執行時，常見的觀察結果是訓練步驟之間存在間隙。在以下追蹤檢視畫面的圖片中，步驟 8 和步驟 9 之間存在很大的間隙，這表示 GPU 在該時間內處於閒置狀態。

如果您的追蹤檢視器顯示步驟之間存在很大的間隙，這可能表示您的程式受輸入限制。在這種情況下，如果您尚未這樣做，則應參閱上一節關於偵錯輸入管線的內容。

不過，即使使用最佳化的輸入管線，由於 CPU 執行緒爭用，步驟結束與另一個步驟開始之間仍可能存在間隙。tf.data 會使用背景執行緒來平行化管線處理。這些執行緒可能會干擾每個步驟開始時發生的 GPU 主機端活動，例如複製資料或排程 GPU 運算。

如果您注意到主機端存在很大的間隙 (主機端會在 GPU 上排程這些運算)，則可以設定環境變數 TF_GPU_THREAD_MODE=gpu_private。這可確保 GPU 核心是從其自己的專用執行緒啟動，且不會在 tf.data 工作之後排隊。

步驟之間的間隙也可能是由指標計算、Keras 回呼或在主機上執行的 tf.function 外部運算所造成。這些運算的效能不如 TensorFlow 圖形內的運算。此外，其中一些運算在 CPU 上執行，並在 GPU 之間來回複製張量。

如果在最佳化輸入管線後，您仍然在追蹤檢視器中注意到步驟之間的間隙，則應查看步驟之間的模型程式碼，並檢查停用回呼/指標是否可改善效能。這些運算的一些詳細資訊也位於追蹤檢視器中 (裝置端和主機端)。在這種情況下，建議的做法是攤銷這些運算的額外負擔，方法是在固定數量的步驟之後 (而不是每個步驟之後) 執行這些運算。當在 tf.keras API 中使用 Model.compile 方法時，設定 steps_per_execution 旗標會自動執行此操作。對於自訂訓練迴圈，請使用 tf.while_loop。

2. 達到更高的裝置使用率

1. 小型 GPU 核心和主機核心啟動延遲

主機會將要 GPU 上執行的核心排入佇列，但在核心實際在 GPU 上執行之前，會產生延遲 (約 20-40 微秒)。在理想情況下，主機會在 GPU 上排入足夠的核心，以便 GPU 將大部分時間用於執行，而不是等待主機排入更多核心。

TensorBoard 上 Profiler 的「總覽頁面」顯示 GPU 因等待主機啟動核心而閒置的時間量。在下圖中，GPU 約有 10% 的步驟時間處於閒置狀態，等待核心啟動。

同一程式的追蹤檢視器顯示核心之間存在小間隙，其中主機正忙於在 GPU 上啟動核心。

透過在 GPU 上啟動許多小型運算 (例如純量加法)，主機可能無法跟上 GPU。TensorBoard 中的 TensorFlow 統計資料工具針對同一設定檔顯示，126,224 個 Mul 運算耗時 2.77 秒。因此，每個核心約為 21.9 微秒，這非常小 (與啟動延遲時間相近)，並可能導致主機核心啟動延遲。

如果您的追蹤檢視器顯示 GPU 上運算之間存在許多小間隙 (如上圖所示)，您可以：

串連小型張量並使用向量化運算，或使用更大的批次大小，使每個啟動的核心執行更多工作，這將使 GPU 保持更長時間的忙碌狀態。
確保您使用 tf.function 建立 TensorFlow 圖形，以便您不會在純粹的渴望模式下執行運算。如果您使用 Model.fit (而不是使用 tf.GradientTape 的自訂訓練迴圈)，則 tf.keras.Model.compile 會自動為您執行此操作。
使用 XLA 和 tf.function(jit_compile=True) 或自動叢集融合核心。如需更多詳細資訊，請前往下方的「3. 啟用混合精度和 XLA」一節，以瞭解如何啟用 XLA 以獲得更高的效能。此功能可以提高裝置使用率。

2. TensorFlow 運算放置

Profiler「總覽頁面」會顯示放置在主機與裝置上的運算百分比 (您也可以透過查看追蹤檢視器來驗證特定運算的放置)。如下圖所示，您會希望主機上的運算百分比與裝置相比非常小。

理想情況下，大多數運算密集型運算都應放置在 GPU 上。

若要找出模型中的運算和張量指派給哪些裝置，請將 tf.debugging.set_log_device_placement(True) 設定為程式的第一個陳述式。

請注意，在某些情況下，即使您指定要將運算放置在特定裝置上，其執行也可能會覆寫此條件 (範例：tf.unique)。即使是單一 GPU 訓練，指定分散式策略 (例如 tf.distribute.OneDeviceStrategy) 也可以讓運算更確定性地放置在您的裝置上。

將大多數運算放置在 GPU 上的原因之一是為了防止主機和裝置之間過多的記憶體複製 (預期主機和裝置之間會進行模型輸入/輸出資料的記憶體複製)。以下追蹤檢視畫面中的 GPU 串流「#167」、「#168」和「#169」示範了過度複製的範例。

如果這些複製作業阻止 GPU 核心執行，則有時可能會損害效能。追蹤檢視器中的記憶體複製運算具有關於作為這些複製張量來源的運算的更多資訊，但將記憶體複製與運算建立關聯可能並非總是容易。在這些情況下，查看附近的運算以檢查記憶體複製是否在每個步驟中的相同位置發生會很有幫助。

3. GPU 上更有效率的核心

一旦您程式的 GPU 使用率可接受，下一步是研究如何透過利用 Tensor Core 或融合運算來提高 GPU 核心的效率。

1. 利用 Tensor Core

新式 NVIDIA® GPU 具有專門的 Tensor Core，可大幅提升符合條件的核心效能。

您可以使用 TensorBoard 的「GPU 核心統計資料」來視覺化哪些 GPU 核心符合 Tensor Core 條件，以及哪些核心正在使用 Tensor Core。啟用 fp16 (請參閱下方的「啟用混合精度」一節) 是使您程式的通用矩陣乘法 (GEMM) 核心 (matmul 運算) 利用 Tensor Core 的一種方法。當精度為 fp16 且輸入/輸出張量維度可被 8 或 16 整除 (int8) 時，GPU 核心會有效率地使用 Tensor Core。

如需關於如何使核心對 GPU 有效率的其他詳細建議，請參閱 NVIDIA® 深度學習效能指南。

2. 融合運算

使用 tf.function(jit_compile=True) 融合較小的運算以形成更大的核心，從而顯著提高效能。若要瞭解更多資訊，請參閱 XLA 指南。

3. 啟用混合精度和 XLA

在遵循上述步驟之後，啟用混合精度和 XLA 是您可以採取的兩個選用步驟，以進一步提高效能。建議的方法是逐一啟用它們，並驗證效能優勢是否如預期。

1. 啟用混合精度

TensorFlow「混合精度」指南說明如何在 GPU 上啟用 fp16 精度。在 NVIDIA® GPU 上啟用 AMP 以使用 Tensor Core，並在 Volta 和更新的 GPU 架構上與僅使用 fp32 (float32) 精度相比，實現高達 3 倍的整體速度提升。

確保矩陣/張量維度滿足呼叫使用 Tensor Core 的核心的需求。當精度為 fp16 且輸入/輸出維度可被 8 或 16 整除 (int8) 時，GPU 核心會有效率地使用 Tensor Core。

請注意，使用 cuDNN v7.6.3 及更高版本，卷積維度將在必要時自動填充，以利用 Tensor Core。

請遵循以下最佳做法，以最大化 fp16 精度的效能優勢。

1. 使用最佳 fp16 核心

啟用 fp16 後，您程式的矩陣乘法 (GEMM) 核心應使用利用 Tensor Core 的對應 fp16 版本。但是，在某些情況下，這不會發生，並且您不會從啟用 fp16 獲得預期的速度提升，因為您的程式會改為回退到效率低下的實作。

「GPU 核心」統計資料頁面顯示哪些運算符合 Tensor Core 條件，以及哪些核心實際正在使用有效率的 Tensor Core。NVIDIA® 深度學習效能指南包含關於如何利用 Tensor Core 的其他建議。此外，使用 fp16 的優勢也將在先前受記憶體限制的核心中顯示出來，因為現在運算將花費一半的時間。

2. 動態與靜態損失縮放

使用 fp16 時，損失縮放是必要的，以防止由於低精度而導致的下溢。損失縮放有兩種類型：動態和靜態，兩者都在「混合精度指南」中更詳細地說明。您可以使用 mixed_float16 策略在 Keras 最佳化工具中自動啟用損失縮放。

當嘗試最佳化效能時，請務必記住，動態損失縮放可能會引入在主機上執行的其他條件運算，並導致追蹤檢視器中步驟之間可見的間隙。另一方面，靜態損失縮放沒有此類額外負擔，並且在效能方面可能是更好的選擇，但需要指定正確的靜態損失縮放值。

2. 使用 tf.function(jit_compile=True) 或自動叢集啟用 XLA

作為使用單一 GPU 獲得最佳效能的最後一步，您可以嘗試啟用 XLA，這將融合運算並提高裝置使用率和降低記憶體佔用量。如需關於如何使用 tf.function(jit_compile=True) 或自動叢集在程式中啟用 XLA 的詳細資訊，請參閱 XLA 指南。

您可以將全域 JIT 層級設定為 -1 (關閉)、1 或 2。較高的層級更積極，可能會降低平行處理並使用更多記憶體。如果您有記憶體限制，請將值設定為 1。請注意，XLA 對於具有可變輸入張量形狀的模型效能不佳，因為 XLA 編譯器每次遇到新形狀都必須保持編譯核心。

2. 優化多 GPU 單一主機的效能

tf.distribute.MirroredStrategy API 可用於將模型訓練從單一 GPU 擴展到單一主機上的多個 GPU。(若要瞭解關於如何使用 TensorFlow 執行分散式訓練的更多資訊，請參閱「使用 TensorFlow 進行分散式訓練」、「使用 GPU」和「使用 TPU」指南，以及「使用 Keras 進行分散式訓練」教學課程。)

雖然從單一 GPU 過渡到多個 GPU 在理想情況下應可立即擴展，但您有時可能會遇到效能問題。

從使用單一 GPU 訓練移至在同一主機上使用多個 GPU 訓練時，在理想情況下，您應該體驗到效能擴展，且僅具有梯度通訊和增加的主機執行緒使用率的額外負擔。由於此額外負擔，如果您從 1 個 GPU 移至 2 個 GPU，則不會獲得精確的 2 倍速度提升。

以下追蹤檢視畫面顯示在多個 GPU 上訓練時額外通訊額外負擔的範例。在執行權重更新之前，串連梯度、跨複本傳輸梯度和分割梯度會產生一些額外負擔。

以下檢查清單將協助您在最佳化多 GPU 情境中的效能時獲得更佳效能：

嘗試最大化批次大小，這將提高裝置使用率並攤銷跨多個 GPU 的通訊成本。使用記憶體分析器有助於瞭解您的程式與峰值記憶體使用率的接近程度。請注意，雖然較大的批次大小可能會影響收斂，但這通常會被效能優勢所抵銷。
從單一 GPU 移至多個 GPU 時，現在同一主機必須處理更多的輸入資料。因此，在 (1) 之後，建議重新檢查輸入管線效能，並確保其不是瓶頸。
檢查程式追蹤檢視畫面中的 GPU 時間軸是否有任何不必要的 AllReduce 呼叫，因為這會導致跨所有裝置的同步處理。在上面顯示的追蹤檢視畫面中，AllReduce 是透過 NCCL 核心完成的，並且每個 GPU 在每個步驟的梯度上只有一個 NCCL 呼叫。
檢查是否有可以最小化的不必要的 D2H、H2D 和 D2D 複製運算。
檢查步驟時間以確保每個複本都執行相同的工作。例如，可能會發生一個 GPU (通常是 GPU0) 超額訂閱，因為主機錯誤地最終在其上放置了更多工作。
最後，檢查追蹤檢視畫面中所有 GPU 的訓練步驟是否有任何循序執行的運算。當您的程式包含從一個 GPU 到另一個 GPU 的控制依賴項時，通常會發生這種情況。過去，在這種情況下偵錯效能是根據具體情況解決的。如果您在程式中觀察到此行為，請提交 GitHub 問題，並附上追蹤檢視畫面的圖片。

1. 優化梯度 AllReduce

當使用同步策略進行訓練時，每個裝置都會接收一部分輸入資料。

在計算模型的前向和反向傳遞之後，需要在每個裝置上計算的梯度進行彙總和縮減。此「梯度 AllReduce」會在每個裝置上的梯度計算之後以及最佳化工具更新模型權重之前發生。

每個 GPU 首先串連跨模型層的梯度，使用 tf.distribute.CrossDeviceOps (預設為 tf.distribute.NcclAllReduce) 跨 GPU 傳輸梯度，然後傳回每個層縮減後的梯度。

最佳化工具將使用這些縮減後的梯度來更新模型的權重。理想情況下，此程序應在所有 GPU 上同時發生，以防止任何額外負擔。

AllReduce 的時間應大致與…

(number of parameters * 4bytes)/ (communication bandwidth)

這個計算很有用，可以快速檢查您在執行分散式訓練任務時的效能是否符合預期，或者是否需要進一步進行效能偵錯。您可以從 Model.summary 取得模型中的參數數量。

請注意，由於 TensorFlow 使用 fp32 (float32) 來傳輸梯度，因此每個模型參數的大小為 4 個位元組。即使您啟用 fp16，NCCL AllReduce 仍會使用 fp32 參數。

為了獲得擴展的優勢，步進時間需要遠高於這些額外負擔。實現此目的的一種方法是使用更大的批次大小，因為批次大小會影響步進時間，但不會影響通訊額外負擔。

2. GPU 主機執行緒競爭

當執行多個 GPU 時，CPU 的工作是透過在裝置之間有效率地啟動 GPU 核心，讓所有裝置保持忙碌。

然而，當有許多獨立的操作 CPU 可以在一個 GPU 上排程時，CPU 可能會決定使用大量的主機執行緒來保持一個 GPU 忙碌，然後以非決定性的順序在另一個 GPU 上啟動核心。這可能會導致傾斜或負向擴展，進而對效能產生負面影響。

下方的追蹤檢視器顯示了當 CPU 無效率地錯開 GPU 核心啟動時的額外負擔，因為 GPU1 處於閒置狀態，然後在 GPU2 啟動後才開始執行運算。

主機的追蹤檢視顯示，主機在 GPU1 上啟動核心之前，先在 GPU2 上啟動核心 (請注意，下方的 tf_Compute* 運算並非 CPU 執行緒的指標)。

如果您在程式的追蹤檢視中遇到這種 GPU 核心錯開的情況，建議的動作是

將 TensorFlow 環境變數 TF_GPU_THREAD_MODE 設定為 gpu_private。此環境變數會告知主機將執行緒保留為 GPU 私有。
預設情況下，TF_GPU_THREAD_MODE=gpu_private 會將執行緒數量設定為 2，這在大多數情況下已足夠。但是，可以透過將 TensorFlow 環境變數 TF_GPU_THREAD_COUNT 設定為所需的執行緒數量來變更該數字。