Folia
← 返回头版

Google开源TorchTPU项目 为PyTorch提供原生TPU后端支持

Google已开源TorchTPU项目,为PyTorch框架提供了原生的张量处理单元(TPU)后端支持1。该项目包含PyTorch ATen内核实现、编译缓存优化和分布式训练支持等功能,使用户能够在Google的TPU上直接运行PyTorch模型和操作1。

TorchTPU项目当前版本为0.2.0,已在PyPI上发布,支持Python 3.11及更高版本1。项目内置了多种分布式训练能力,包括数据并行和张量并行两种模式1,并与torch.compile()集成以支持图模式执行1。官方提供的示例模型涵盖Llama、Qwen、ResNet、DLRM、minGPT和UNet等多个应用场景1。

该项目采用Apache 2.0许可证,使用Bazel进行依赖管理和构建,欢迎社区贡献1。


评论