Accelerate 文档

TPU 训练

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始入门

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()
< > 在 GitHub 上更新