在 TPU 節點上訓練
TPU 的設計目的是擴展至 TPU Pod。TPU Pod 是由專用高速網路介面連結的 TPU 裝置集合。TPU Pod 可讓您將處理負載分散至多個 TPU。每塊 TPU 板都連線至高效能的 CPU 型主機,用於載入和預先處理資料等作業。如要充分運用更多 TPU,您必須微調多個訓練工作參數。
使用 TPU Pod 進行訓練的設定方式會因架構而異。 請使用下列連結,查看各架構的 Pod 訓練詳細資訊:
下列各節將說明一些常見問題、模型中需要進行的變更,以及減少或避免 Pod 故障的最佳做法。
針對批量訓練次數進行資源調度
如要針對較大的 TPU 類型執行線性資源調度,請維持相同的每核心批量。
舉例來說,若您在 v6e-8 上使用 1024 的批量,請在 v6e-32 上使用 4096 (4 * 1024) 的批量。這樣即可充分運用 TPU 硬體。您可以使用較小的批次大小,但如此一來訓練將無法進行線性資源調度。
部分模型包含 train_steps 旗標,其中一個步驟對應處理單一資料批次。增加批次大小時,請縮減訓練步驟數量,讓訓練範例總數保持不變。
舉例來說,如果 100 個步驟的批次大小為 1000,訓練期間會處理 100,000 個範例。如果現在有 4 個工作站,有效批次大小為 4000,則必須將步數調整為 25,才能處理相同的 100,000 個範例。如果模型使用 epochs 旗標,則無須調整訓練批次數。
較大的批量可變更模型的收斂行為,因此您也可以微調一些超參數,例如學習率。
在與 TPU Pod 相同的地區中,使用區域 Cloud Storage bucket
一般而言,TPU 訓練的最佳做法是一律使用同一地區的資源。使用 TPU Pod 時,資源地區尤為重要,因為 Cloud Storage 值區和 TPU 位於同一地區時,資料移轉率會更高。
確認在與 TPU 相同的地區中,使用區域 Cloud Storage 值區來訓練資料集和檢查點。
在 TPU Pod 上開發時的工作流程最佳做法
開發新的 TPU 工作負載時,通常最好先從最小的 TPU 開始開發,然後逐步疊代至較大的 TPU 大小。請先使用小型 TPU 版本 (例如 v6e-8)。
- 測試工作負載是否符合預期行為
- 使用效能工具測試及驗證效能
工作負載運作正常並達到效能目標後,請擴充至較大的 TPU 類型,例如 v6e-32。逐步且反覆增加 TPU 大小,同時驗證擴充性 (功能和效能),直到達到所需的 TPU 大小為止。