使用 bfloat16 提升模型效能

根據預設,TPU 會使用 bfloat16 值執行矩陣相乘作業,並使用 IEEE float32 值進行累加。使用精度較低的浮點數,可縮短收斂時間,且不會影響準確度。

bfloat16float32 的動態範圍相同。不過,bfloat16 會使用一半的記憶體空間。如要進一步瞭解 bfloat16 效能,請參閱「A Study of BFLOAT16 for Deep Learning Training」。

明確使用 bfloat16

雖然 TPU 中的自動格式轉換可讓您不必考慮數值精確度,但您可以明確將值轉換為 bfloat16,藉此提升效能。明確將值轉換為 bfloat16 有兩個原因:

  1. bfloat16 格式儲存值可節省晶片上記憶體,讓 Cloud TPU 訓練更大的模型或使用更大的批次大小。

  2. 某些作業受限於記憶體頻寬,這表示從記憶體載入資料所需的時間可能會拖慢整體運算作業所需的時間。以 bfloat16 格式儲存這些運算子的運算元和輸出,可減少必須轉移的資料量,進而改善整體速度。

建議您先試試 Cloud TPU 參考模型,接著,分析工具指南疑難排解指南會提供深入的技術資訊,協助您自行建立及最佳化機器學習模型。

格式轉換詳細資料

XLA 編譯器會自動插入從 float32bfloat16 的格式轉換。在 TPU 上,轉換中的捨入方式為「四捨五入」,並溢位至 inf。此外,Cloud TPU 上的 bfloat16 不支援次常態值,因此所有次常態值會在轉換期間清除為零。特殊值 (例如 NaNinf) 會在轉換過程中保留。

XLA 編譯器也會自動插入從 bfloat16float32 的格式轉換。由於 float32 可代表 bfloat16 中的所有確切值,因此轉換會在尾數位元中填入 16 個零。特殊值會在轉換中保留。

從在 Cloud TPU 上訓練的模型取得的檢查點,可部署至其他硬體平台 (例如在 CPU 或 GPU 上進行推論或微調),無須進行大量手動轉換。