在 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
旗標,其中一個步驟會對應處理單一資料批次。增加批量時,請縮減訓練步驟的數量,讓訓練範例的總數保持不變。
舉例來說,如果批次大小為 1000,訓練步驟為 100 步,則訓練期間會處理 100,000 個範例。如果您現在有 4 個 worker,且有效批量大小為 4000,則必須將步驟數調整為 25,才能處理相同的 100,000 個範例。如果模型使用 epochs
標記,則無須調整步驟數。
較大的批量可變更模型的收斂行為,因此您也可以微調一些超參數,例如學習率。
在與 TPU Pod 相同的地區中,使用區域 Cloud Storage 值區
一般而言,TPU 訓練的最佳做法是一律使用同一地區的資源。使用 TPU Pod 時,資源地區尤為重要,因為 Cloud Storage 值區和 TPU 位於同一地區時,資料移轉率會更高。
確認在與 TPU 相同的地區中,使用區域 Cloud Storage 值區來訓練資料集和檢查點。
在 TPU Pod 上開發時的最佳工作流程
開發新的 TPU 工作負載時,通常最好從最小的 TPU 開始開發,然後逐步迭代至較大的 TPU 大小。請先使用小型 TPU 版本 (例如 v6e-8)。
- 測試工作負載是否符合預期行為
- 使用成效工具測試及驗證成效
工作負載正常運作並達到效能目標後,請擴展至較大的 TPU 類型,例如 v6e-32。逐步且重複增加 TPU 大小,同時驗證可擴充性 (功能和效能),直到達到所需的 TPU 大小為止。