TensorFlow - TPU 疑難排解
這份指南與常見問題一併提供疑難排解說明,協助使用者在 Cloud TPU 上訓練 TensorFlow 模型。如果您要排解 PyTorch 或 JAX 訓練問題,可以參考這些架構的疑難排解文件:
如需 Cloud TPU 的一般使用指南,請參閱:
總覽
Cloud TPU 常見問題可分為以下幾類:
連結 TPU 伺服器的問題
本節說明如何排解 TensorFlow 連線至 TPU 時停止回應或顯示錯誤訊息的情況。執行大型模型時,TPU 圖形編譯步驟可能需要較長時間完成,所以至少要在指令碼執行 5 分鐘後,才能斷定是否停止回應。
首先,請檢查究竟是伺服器本身發生問題,還是 TensorFlow 訓練管道出現狀況。如要執行這項操作,請執行 TensorFlow 程式,並確認其運作正常。如果仍有連線問題,表示問題出在 TPU 伺服器。在這種情況下:
執行以下指令,列出可用的 TPU。將 zone 和 project-id 替換為您的區域和專案 ID。
(vm)$ gcloud compute tpus tpu-vm list --zone zone --project project-id
完成後將顯示下列輸出:
NAME ZONE ACCELERATOR_TYPE NETWORK_ENDPOINT NETWORK RANGE STATUS TPU_NAME us-central1-b v2-8 10.240.1.2:8470 default 10.240.1.0 READY
確認該 TPU 已列為
READY
。如果您發現 TPU 並未顯示
READY
,或是仍有連線問題,請使用以下指令手動重新啟動伺服器:(vm)$ gcloud compute tpus tpu-vm stop TPU_NAME && gcloud compute tpus tpu-vm start TPU_NAME
這項作業可能需要幾分鐘才能完成。
重新執行
gcloud compute tpus tpu-vm list
指令,並等候 TPU 進入READY
狀態。這可能需要幾分鐘的時間。請嘗試再次執行程式。
如果仍無法解決問題,請使用取得支援中所述的其中一種機制尋求協助。
如果程式碼可以正常運作,但模型仍停止回應,則可能是訓練管道發生問題。如要對此進行偵錯,請先將程式碼中的 TPUStrategy 替換為預設策略。使用預設策略時,無論您在何處使用 strategy.scope()
或 strategy.run()
,模型都會在 CPU (或 GPU,如有) 上執行,而非在 TPU 上執行。如果模型在 CPU 上執行,而非 TPU,則一定有 TPU 專屬問題。如果仍無法執行,最佳做法是在 CPU 上偵錯。
訓練期間 ssh
連線中斷
在長時間訓練期間,您與 Cloud TPU 的 ssh
連線可能會逾時 (特別是使用 Cloud Shell 時)。此時,TPU 主控台不會顯示任何輸出內容,因此可能會誤以為 TPU 已停止訓練。為避免這種情況,請使用終端機多工器或工作階段管理工具 (例如 tmux
或 screen
) 執行訓練工作階段。無論訓練時間長短為何,這麼做都能讓 ssh
連線保持啟用狀態。
常見錯誤的除錯方式
本節說明如何排解在 Cloud TPU 上訓練模型時可能遇到的常見錯誤。
無法建立 TPU
建立 Cloud TPU 時,您可能會看到下列錯誤:
googleapiclient.errors.HttpError: < HttpError 403 when requesting https://content-tpu.googleapis.com/v1/projects/{PROJECT}/locations/{ZONE}/nodes/{TPU_NAME}?alt=json returned "Request had insufficient authentication scopes."
這是權限問題,您可以執行下列指令來解決:
gcloud auth login --update-adc
這項指令會更新應用程式預設憑證 (ADC),應該可以解決問題。詳情請參閱 gcloud auth login。
不支援動態形狀
錯誤訊息
ValueError: shape [Shape] must have a fixed size for dimension d that is known at graph construction time.
受影響的架構和設定
這則訊息只會在使用 TensorFlow 進行 XLA 編譯時出現。
詳細資料
在 TPU 上執行模型時,Cloud TPU 會使用 XLA 編譯器編譯模型。雖然這個編譯步驟可大幅改善訓練速度和記憶體用量,但圖中所有張量的形狀 (維度大小) 都必須是靜態的,也就是說,圖形編譯時必須能夠辨認這些張量的值。如果編譯時無法判定任一形狀,TPU 編譯就會失敗,並顯示類似上述範例的錯誤訊息。
dataset.batch(batch_size)
是傳回動態形狀時常用的運算。由於串流中剩餘的樣本數可能少於批次大小,因此在 TPU 上進行訓練時,請為 dataset.batch
設定 drop remainder=True
。這個運算應該會捨棄檔案中的最後幾個樣本,進而確保每個批次都有 batch_size 的靜態形狀。例如:
dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)
無法使用的 TensorFlow 運算
錯誤訊息
NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT devices compatible with node
受影響的架構和設定
使用 TensorFlow 訓練時,可能會顯示這則訊息。
詳細資料
模型所用的 TensorFlow 運算無法在 TPU 上執行。
如需可在 TPU 上使用的運算清單、日後支援資源和建議解決方案,請參閱本指南中的可用 TensorFlow 運算。
記憶體不足的錯誤訊息
錯誤訊息
ResourceExhaustedError: Ran out of memory in memory space hbm; used: YYY; limit: 7.48G.
受影響的架構和設定
使用 TensorFlow、PyTorch 或 JAX 訓練時,可能會出現這則訊息。
詳細資料
每個 Cloud TPU 都是由八個 TPU 核心組成,v2 TPU 有 8 GB,v3 TPU 則有 16 GB 的 RAM (或 HBM,即高頻寬記憶體)。這個記憶體用於儲存權重 (變數) 張量,以及梯度計算所需的中繼結果張量。如果模型太大,無法放入 TPU RAM,初始化作業就會失敗,並顯示錯誤訊息。如需詳細說明,請參閱減少記憶體用量一節。
減少記憶體用量的提示:
- 檢查是否有過多的張量填充
- 使用 bfloat16 格式
- 如果輸入大小或模型過大,您可以使用 TensorFlow 的實驗模型平行處理功能來解決模型大小問題。
停止執行時發生問題
如果 TensorFlow 在 TPU 執行期間發生錯誤,指令碼有時可能會停止回應,而非自行關閉並前往殼層。如果發生這種情況,請按下鍵盤上的 CTRL+C
以觸發 SIGQUIT
,讓 Python 程式立即結束。
同樣地,在 TPU 執行期間按下 CTRL+C
並「不」會立即關閉 TensorFlow,而是會等到目前的疊代迴圈結束才會完全關閉。
透過這個方法結束程式後,如果在重新連線至 TPU 時發生如 DeadlineExceededError 這類新錯誤,請使用下列指令手動重設 TPU 伺服器:
gcloud compute tpus tpu-vm stop tpu-name --zone=zone gcloud compute tpus tpu-vm start tpu-name --zone=zone
其中 tpu-name 是從 gcloud compute tpus tpu-vm list
指令顯示的第一欄取得,而 zone 是第二欄顯示的區域。
過度的張量填充
記憶體發生問題的可能原因
TPU 會填充記憶體中的張量,也就是說,TPU 會增加儲存在記憶體中的張量大小,以便提升執行運算時的效率。填充過程會不著痕跡地在硬體層級間進行,且不會影響結果。但在某些情況下,填充將造成記憶體用量和執行時間大幅增加。
如何減少記憶體用量
TPU 軟體會試圖配置記憶體中的張量,藉此增加運算效率並降低填充量。雖然這種記憶體的配置程序相當複雜,但為了取得最佳結果,仍請依照以下準則執行模型。如果要將記憶體負擔降至最低,並同時獲得最高的運算效率,就必須滿足下列其中一項條件:
總批量應為 64 的倍數 (每個 TPU 核心為 8),且特徵維度應為 128 的倍數
或
總批量應為 1024 的倍數 (每個 TPU 核心為 128),且特徵維度應為 8 的倍數
批量為 1024「且」特徵維度為 128 的倍數時,將可獲得最佳效率。不過這樣的配置可能不適用於所有模型。在此詳加說明,「特徵維度」是指全連結層的隱藏大小或卷積層中的輸出通道數,並非所有層都符合這個規則,特別是網路的第一層和最後一層。不過這樣也無妨,因為大部分的模型本來就需要一些填充量。
減少記憶體用量
如果您在 TPU 上執行模型時遇到記憶體不足錯誤,就必須採取步驟來減少模型的記憶體用量。
減少記憶體用量的最有效方法如下:
- 減少過度的張量填充
- 減少批次大小
批次大小或模型過大
記憶體發生問題的可能原因
在 CPU、GPU 或 TPU 上訓練神經網路時,記憶體用量主要來自下列兩個地方:
- 記憶體用量與模型中的權重數成正比。
- 儲存計算反向傳播所需的正向傳播中繼啟動。記憶體用量與批量、層大小和層數成正比。
由此可知,模型所需的記憶體主要取決於批量。
模型所需的記憶體用量取決於網路中的層數。
TPU 執行階段會嘗試最佳化運算子,以配合記憶體中的模型大小 (這個動作稱為再實質化,類似建立梯度查核點),但不一定能成功。
如何減少記憶體用量
請逐漸縮減批次大小,直到符合記憶體容量為止,另外請同時確認總批次大小是 64 的倍數 (每核心批次大小須為 8 的倍數)。提醒您,TPU 上的批次大小越大,就越能提升效率。一開始的時候,建議將總批次大小設為 1024 (每核心為 128)。
如果使用小規模的批量 (例如 64),還是無法在 TPU 上執行模型,請試著減少層數或層大小。
提升訓練速度
如果您的模型能在 TPU 上順利執行,但訓練速度低於預期,請參閱本節查看提升速度的可行方法。如要瞭解其他改善訓練成效的建議,請參閱成效指南。
每個訓練迴圈的執行步驟太少
效能問題說明
將 steps_per_execution
引數傳遞至 Model.compile
,即可控制在主機回呼之間執行多少訓練步驟。每個主機回呼都需要在 TPU 伺服器的主機 CPU 和 TPU 裝置之間進行大量通訊,因此如果 steps_per_execution
太小,訓練速度就會降低。
如何判斷您的模型是否受到影響
如果 TPU 設定檔顯示 TPU 裝置步驟之間經常發生主機 CPU 回呼,則訓練可從較大的 steps_per_execution
值中受益。
解決方式
將 steps_per_execution
設為更大的值。請注意,steps_per_execution
可設為很大的值,但請記住,記錄訊息和儲存查核點只能在執行指定步驟數量後發生。
輸入處理瓶頸
效能問題說明
當 TPU 在訓練特定資料區塊時,輸入處理函式會在 CPU 上準備下一個資料區塊。如果輸入函式的執行時間比模型函式「長」,則在輸入函式擷取資料時,TPU 會處於閒置狀態。
如何判斷您的模型是否受到影響
請按照 Cloud TPU 工具:輸入管道分析工具中的操作說明,查看 TensorBoard 中的輸入管道分析:
輸入管道分析頁面會顯示清楚的摘要,說明您的模型是否因輸入處理而出現瓶頸,同一個頁面也會顯示每個作業的執行時間,讓您找出有問題的作業。
解決方式
透過 Dataset
API 載入資料時,可採用下列解決方式:
- 將資料以
tf.train.Example
結構集合的形式儲存在TFRecord
檔案中,然後使用TFRecordDataset
載入這些資料。如需相關範例,請參閱 Dataset API 教學課程或 ResNet 教學課程。 - 使用
dataset.cache()
或dataset.prefetch()
將輸入資料置於緩衝區。這樣一來,就能避免檔案存取速度變慢的偶發狀況,進而防止瓶頸形成。 - 指定
dataset.map()
函式的num_parallel_calls
參數,藉此啟用多執行緒map()
作業。num_parallel_calls
值的啟發式是使用可用的 CPU 核心數量。 - 針對佔用大量資源的資料,進行離線的前置處理作業。這樣一來,就只會產生一次性費用,而非在每個訓練的各訓練週期中重複產生費用。
所有輸入處理都在位於 TPU 伺服器的 CPU 中執行,而「不是」在本機電腦上執行,因此本機電腦的速度並非決定性因素。
步驟耗時過長且 MXU 使用率偏低
效能問題說明
Cloud TPU 能以極高速執行矩陣乘積和卷積運算。雖然 TPU 也能有效執行其他多項 TensorFlow 運算,但若與其他硬體相比,這些運算並非 TPU 的主要優勢。因此,建議使用以矩陣乘積或卷積運算為主的模型,充分發揮 TPU 的效能。
如何判斷您的模型是否受到影響
在這種情況下,您會看到的症狀是步驟時間緩慢,且剖析效能時顯示 MXU 使用率偏低。
解決方式
請盡量減少非矩陣相乘運算作業的數量。減少矩陣乘積的次數後,請重新進行基準測試,看看 TPU 的效能是否可接受。
過度的張量填充
效能問題說明
TPU 會填充記憶體中的張量,以便有效利用運算單元。填充會增加記憶體和記憶體頻寬的用量。如需進一步瞭解張量填充,或修正相關問題,請參閱張量填充一節。
處理量緩慢且記憶體用量偏低
效能問題說明
一般來說,使用的批次大小越大,模型在 TPU 上的訓練速度也就越快 (以每秒樣本數為單位)。
如何判斷您的模型是否受到影響
不論模型類別,都應將批次大小設為 64 以上 (每個 TPU 核心為 8),因為 TPU 一定會將張量填充到這個大小。在 TPU 上進行訓練時,理想的批次大小是 1024 (每個 TPU 核心為 128)。這樣一來,就不會產生與記憶體和填充相關的效率問題。
解決方式
最佳做法是使用記憶體可容納的最大批次大小,且為 64 的倍數。最簡單的方式就是在一開始的時候,將批次大小設為 1024。如果發生記憶體不足的錯誤,可以試著縮減批次大小,直到模型能順利執行為止。變更模型的批次大小時,可能還需要調整其他超參數,才能達到相同的模型準確度 (例如學習率),但仍須視個案進行評估。
層的大小太小
效能問題說明
即使模型以矩陣乘積或卷積做為主要運算方式,如果輸入張量太小,還是無法以最高效率執行 TPU。和其他硬體相比,當批次大小和層的大小都大時 (例如,維度 >= 512),TPU 的執行效率最高。
如何判斷您的模型是否受到影響
一般來說,使用低於 128 的層大小,將導致 TPU 效率不彰,這是因為 128 是 TPU 矩陣乘積單元的內建維度。如果是全連結層,建議將最低隱藏大小設為 512,藉此提升效率。請注意,卷積層通常不需和全連結層一樣大,即可達到相同的效率。
解決方式
如果您希望模型的層大小較小,主要原因是為了提升訓練速度,請在 TPU 上使用較大的層重新測試模型。舉例來說,即使模型執行的運算量是原來的 2 倍,將某個層的輸出大小從 256 增加到 512,訓練時間可能只會增加 20%。
運算層模型剖析
測量運算層的執行時間和記憶體用量,通常是找出效能瓶頸的好方法。如需相關操作說明:
請參閱 Cloud TPU 工具:追蹤記錄檢視器。
模型準確度降低的除錯方法
Cloud TPU 生態系統的其中一項目標,就是要讓目前在 CPU 或 GPU 上訓練的所有模型,在微調幾項超參數 (例如批次大小和學習率) 後,就能在 TPU 的訓練中達到同等的準確度。然而,在 TPU 上訓練模型時,使用者可能會不免發現準確度有所降低。由於類神經網路訓練的隨機性質,針對這類問題進行除錯可能會是艱鉅的考驗。本節內容說明將模型移植至 TPU 後,如何找出導致模型準確度下降的根本原因。
瞭解資料分割 (資料平行處理)
TensorFlow 的主要目標之一,就是要讓各項運算無論在 CPU、GPU,還是 TPU 上都能產生幾乎相同的結果(隨機運算等特例除外)。一般來說,只要您發現 TPU 和 CPU 上的非隨機運算結果出現極大的差異,都可以回報錯誤。
然而,「單」就訓練管道而言,在 CPU/GPU 和 TPU 上進行的訓練原本就會產生極大的差異,在 TPU 上訓練時,TensorFlow 會執行資料分割,每個 Cloud TPU 都包含 8 個 TPU 核心,這些核心會以獨立的處理單元運作。在訓練的各個步驟中,每個 TPU 核心都會接收一批資料、計算權重梯度、與其他 TPU 核心交換梯度,然後計算權重更新值。根據預設,損失值是採所有核心的平均,但也可以變更 CrossShardOptimizer
的參數,改採總和。
如果可按每獨立樣本的損失平均 (或總和) 計算出模型的總損失,那麼這項程序在數學上就等同於訓練一個單一大型批次。
「不是」按每獨立樣本執行的運算中,最常見的就是批次標準化,此運算會針對批次中的各個核心分別執行。舉例來說,如果總批次大小為 128,則每個核心的批次大小為 16,而每個核心會針對自己的 16 個樣本執行批次規格化。在某些情況下,在小批次 (例如 32 以下) 上執行批次標準化會降低準確度。在理想情況下,總批量應為較大 (例如 256 到 1024)。如果批次大小太大,記憶體無法容納,就必須視個案評估資料分割的效果。
對多核心 TPU 訓練進行除錯
如果模型在 CPU 和單核心 TPU 上得到的損失值「確實」相同,那麼問題原因可能是下列其中一項:
(a) 使用不同的初始化訓練神經模型時,因為產生自然隨機變異數,而導致準確度下降。
(b) 由於 TPU 上資料分割的相關問題,而導致準確度下降。
如要判斷 (a) 是否為問題原因,請使用相同的權重初始化,在 CPU/GPU 和多核心 TPU 上重新訓練整個模型。
如果您認為準確度下降會對統計資料產生嚴重影響,下列是與資料分割相關的問題中,最有可能發生的幾項:
- 如果模型使用批次標準化,只要總批量低於 256 (例如,每個核心低於 32),準確度就有可能降低。
- 批次損失函式會受到資料分割影響。這類損失函式通常用於高度專業的領域。例如,Karras 等人 2017在訓練生成對抗網路時,就曾使用批次鑑別器。
gcloud
設定疑難排解
- 問題
gcloud components update
會顯示以下錯誤訊息:
ERROR: (gcloud.components.update) You cannot perform this action because the Cloud SDK component manager is disabled for this installation.
- 解決方案
如要使用 Google Cloud CLI,您必須使用非透過套件管理員管理的安裝程序。
執行下列指令,移除目前的 gcloud CLI 安裝:
sudo apt-get remove google-cloud-sdk
按照操作說明安裝 Google Cloud CLI。
- 問題
gcloud compute tpus tpu-vm ssh TPU_NAME --zone ZONE
指令會顯示以下錯誤訊息:Waiting for SSH key to propagate. ssh: connect to host 34.91.136.59 port 22: Connection timed out ssh: connect to host 34.91.136.59 port 22: Connection timed out ssh: connect to host 34.91.136.59 port 22: Connection timed out ERROR: (gcloud.compute.tpus.tpu-vm.ssh) Could not SSH into the instance. It is possible that your SSH key has not propagated to the instance yet. Try running this command again. If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.
- 解決方案
安全殼層金鑰傳播作業可能發生錯誤。請嘗試將自動產生的金鑰移至備份位置,強制
gcloud
重新建立金鑰:mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub
偵錯記錄
支援的 Cloud TPU 架構 (JAX、PyTorch 和 TensorFlow) 會使用每個 TPU VM 上提供的共用程式庫 libtpu
存取 TPU。這個程式庫包含用於編譯 TPU 程式的 XLA 編譯器、用於執行已編譯程式的 TPU 執行階段,以及執行階段用於低階存取 TPU 的 TPU 驅動程式。
libtpu
程式庫會記錄有助於偵錯的資訊。根據預設,這些記錄會寫入每個 Cloud TPU VM 的 /tmp/tpu_logs
。您可以在開始訓練前設定下列環境變數,以修改記錄行為:
- TPU_LOG_DIR:記錄檔寫入的目錄
- 目錄位置預設為
/tmp/tpu_logs
。如果目錄尚不存在,系統會建立目錄,但不會建立父目錄。如果在尋找或建立指定目錄時發生錯誤,系統會將訊息輸出至 stderr,但不會停止程式,且記錄功能會停用。將目錄名稱設為「disabled」,即可完全停用寫入磁碟的記錄功能。 - TPU_MIN_LOG_LEVEL:記錄至磁碟的最低嚴重性
- 可選值為 0 (INFO)、1 (WARNING)、2 (ERROR) 和 3 (FATAL)。預設值為 0。
- TPU_STDERR_LOG_LEVEL:除了磁碟外,記錄至 stderr 的最低嚴重性 (如適用)
- 可用的選項與 TPU_MIN_LOG_LEVEL 相同。預設值為 3。
- TPU_MAX_LOG_SIZE_MB:每個記錄檔的大小上限 (以 MB 為單位)
- 當舊的記錄檔大約達到這個大小時,系統會自動啟動新的記錄檔。預設值為 1024。