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。
Google has released TorchTPU as an open-source project, providing native TPU backend support for PyTorch 1. The project enables users to run PyTorch models and operations directly on Google's Tensor Processing Units (TPUs) by offering PyTorch ATen kernel implementations, compilation cache optimizations, and distributed training capabilities 1.
Version 0.2.0 of TorchTPU is now available on PyPI and requires Python 3.11 or later 1. The project supports both data parallelism and tensor parallelism for distributed training scenarios 1. It integrates with torch.compile() to enable graph mode execution 1 and includes reference implementations for several models including Llama, Qwen, ResNet, DLRM, minGPT, and UNet 1. The project is published under the Apache 2.0 license and uses Bazel for dependency management and build processes 1, welcoming community contributions 1.
评论
还没有评论,欢迎留下第一条。