谷歌开源用于TPU上PyTorch的TorchTPU

谷歌已开源TorchTPU，这是一个PyTorch后端，使用户能够在谷歌的张量处理单元（TPU）上运行PyTorch模型和操作。该项目包括针对一系列PyTorch操作的自定义ATen内核、用于优化XLA计算执行的编译缓存，以及分布式训练和模型推理的示例。

主要功能包括通过数据并行和张量并行支持分布式训练、与torch.compile()集成以实现图模式执行，以及用于调试和基准测试的实用工具。示例模型涵盖Llama、Qwen、ResNet、DLRM、minGPT和UNet。

项目提供用户指南和交互式Marimo笔记本，其使用方式类似于Jupyter。笔记本必须在TPU VM上运行，推荐使用v6e GCE实例作为最简单的选项。目前支持Python 3.11及更高版本，与PyTorch的发布兼容性矩阵一致。

最新公开版本0.2.0可通过pip从PyPI安装，并使用来自额外索引的PyTorch CPU版本。开发者也可以使用Bazel从源代码构建wheel包，默认启用优化构建，对于没有凭据的用户，可选择禁用远程构建执行（Remote Build Execution）。