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()