Google 将 TorchTPU 集成进 TorchTitan,发布面向 Cloud TPU 的大模型训练与优化方案
Google 发布 TorchTitan TPU,借助新推出的 TorchTPU 原生后端(torch.device("tpu")),PyTorch 训练脚本只需少量设备初始化改动即可在 Cloud TPU 上训练,分支仓库为 https://github.com/google-pytorch/torchtitan/,由 Google 与 Meta 共同加入 TPU 原生执行。
推荐理由:原文给出按顺序执行的 TPU 训练优化清单,并用三个模型的案例标出每步改动换来的 MFU 或吞吐提升,便于迁移到自己的训练任务。