TPU 訓練

TPU（張量處理單元）是一種專為高效訓練模型而設計的硬體。Accelerate 支援 TPU 訓練，但你需要注意一些事項，即圖編譯。本教程簡要討論編譯，更多詳情請參閱 使用 Accelerate 在 TPU 上訓練 指南。

編譯

TPU 會為訓練步驟中的所有操作（如前向傳播、反向傳播和最佳化器步驟）建立一個計算圖。這就是為什麼第一個訓練步驟總是需要很長時間，因為構建和編譯這個計算圖需要時間。但一旦編譯完成，它就會被快取，所有後續步驟都會快得多。

關鍵是要避免再次編譯你的程式碼，否則訓練會非常慢。這意味著你的所有操作必須完全相同：

• 你的批次中所有張量的長度必須相同（例如，對於 NLP 任務，不能使用動態填充）
• 你的程式碼必須是靜態的（例如，不能有包含 for 迴圈且迴圈長度依賴於輸入的層，比如 LSTM）

權重繫結

一種常見的語言模型設計是繫結嵌入層和 softmax 層的權重。然而，將模型移動到 TPU 上（無論是你自己操作還是將其傳遞給 prepare() 方法）會破壞權重繫結，你需要重新繫結權重。

要在你的指令碼中為 TPU 新增特殊行為（如權重繫結），首先將 `distributed_type` 設定為 `DistributedType.TPU`。然後你可以使用 `tie_weights` 方法來繫結權重。

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()