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