TPU 训练

一个 TPU (Tensor Processing Unit) 是一种专门为高效训练模型而设计的硬件。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()