算力平台:
TPU 训练
TPU (Tensor Processing Unit) 是一种专门设计用于高效训练模型的硬件。Accelerate 支持 TPU 训练,但有一些需要注意的事项,主要是图编译。本教程简要讨论了编译,更多详细信息,请参阅 使用 Accelerate 在 TPUs 上训练 指南。
编译
TPU 会创建一个包含训练步骤中所有操作的图,例如前向传播、反向传播和优化器步骤。这就是为什么第一个训练步骤总是需要一段时间,因为构建和编译这个图需要时间。但一旦编译完成,图会被缓存,所有后续步骤都会快得多。
关键是要避免再次编译代码,否则训练会非常慢。这意味着所有操作必须完全相同:
- 批次中的所有张量必须具有相同的长度(例如,NLP 任务中不能使用动态填充)
- 代码必须是静态的(例如,不能使用根据输入长度不同的循环层,如 LSTM)
权重绑定
常见的语言模型设计是将嵌入层和 softmax 层的权重绑定在一起。然而,将模型移动到 TPU(无论是你自己移动还是通过 [~Accelerator.prepare
] 方法传递)会破坏权重绑定,你需要重新绑定权重。
要在脚本中为 TPU 添加特殊行为(如权重绑定),首先将 [~Accelerator.distributed_type
] 设置为 DistributedType.TPU
。然后你可以使用 [~transformers.PreTrainedModel.tie_weights
] 方法来绑定权重。
py
if accelerator.distributed_type == DistributedType.TPU:
model.tie_weights()