描述 该项目为谷歌的张量处理单元（TPUs）提供了一个PyTorch后端，使用户能够在TPU设备上运行PyTorch模型和操作。它包括用于各种PyTorch操作的自定义ATen内核，优化执行的编译缓存，以及分布式训练和模型推理的示例。 特点 - TPU的PyTorch ATen内核实现。 - 优化执行XLA计算的编译缓存。 - 支持分布式训练（数据并行、张量并行）。 - 各种模型的示例（Llama、Qwen、ResNet、DLRM、minGPT、UNet）。 - 支持与torch.compile()的图模式执行的集成。 - 调试和基准测试的实用工具。 用户指南和互动笔记 用户指南可在https://torchtpu.readthedocs.io/获取。 这是互动笔记的静态版本。它们在docs文件夹中可用。笔记本是Marimo（类似于Jupyter）。您可以通过marimo edit file.py运行它们，然后在浏览器中访问它们。您需要在TPU虚拟机上运行此程序，v6e GCE实例是最简单的。 安装 上述用户指南中有相同的说明。TorchTPU的Python版本支持政策与PyTorch的发布兼容性矩阵相匹配。当前支持Python 3.11及以上版本。有关更多详细信息，请参见Python版本支持政策。 首先创建一个Python虚拟环境： ``` python -m venv .venv source .venv/bin/activate ``` 或者，使用uv： ``` uv venv .venv source .venv/bin/activate ``` 通过pip安装 安装最新的公共版本（在PyPI上的版本为0.2.0）（使用uv时在所有pip命令前加上uv）： ``` pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu ``` 从源代码构建 可以通过以下命令构建Wheels： ``` bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel ``` 所有OSS Bazel构建默认在优化模式（-c opt）下运行，自动通过.bazelrc配置。不优化的构建（-c fastbuild，-c dbg）可能会降低性能，且不用于wheel构建。如果您正在开发或调试C++代码，并明确想要一个未优化的wheel，请使用以下命令： ``` bazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel \ --//:allow_unoptimized_wheel=True ``` 同时要求使用--config=wheel_common，如果没有，它将在构建时失败。该配置设置了--//:wheel_build=True，将共享的XLA/MLIR后端通过torch_version-reset转换路由，使得pywrap将其因素化为单个libxla_base.so；如果不带该标志构建，每个版本的common都会发送其自己的后端副本，导致最终的wheel在导入torch时因为重复的静态注册而中止。wheel_common还开启了远程构建执行（RBE）远程缓存和远程执行器。如果您没有RBE凭证，请附加--config=no_rbe以去除这些标志： ``` bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe ``` 顺序在此很重要：--config=no_rbe必须在--config=wheel_common之后，因为后者配置是胜出的，用于它们都设置的远程缓存和远程执行器标志。通过以下命令安装wheel： ``` cd < path_to_repo > python -m venv .venv ; source .venv/bin/activate # extra-index-url提供了CPU版本的torch，而不是默认的CUDA构建。 pip install ../bazel-bin/ci/wheel/* .whl --extra-index-url https://download.pytorch.org/whl/cpu ``` 注意：如果不带--extra-index-url https://download.pytorch.org/whl/cpu进行安装，则会安装CUDA版本的torch；请根据上面的说明传递--extra-index-url或单独安装CPU构建。 开始使用 以下是一些帮助您入门的指引： - PyTorch：官方PyTorch文档 - 教程和示例：使用TorchTPU的易于理解的PyTorch代码 - 公共API：您可以使用的TorchTPU API 依赖关系 该项目使用Bazel进行依赖管理。主要的Python依赖关系有：待定 Bazel安装 该项目使用Bazel进行构建和依赖管理。要安装Bazel，我们实际上安装Bazelisk，它是一个透明的包装器，可以处理Bazel的下载和安装。安装Bazelisk的方法有几种，这里是一些常见的方法： - 如果您安装了Go： ``` go install github.com/bazelbuild/bazelisk@latest ``` - 直接下载二进制文件： ``` wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64 ``` 请参阅Bazelisk README以获取更多安装说明。无论您如何安装，都请将其位置添加到PATH。 要验证安装，请在repo的根目录中运行以下命令： ``` bazel info ``` 应该会输出关于repo的信息，没有错误。 VSCode clangd设置 对于C++代码导航，建议使用VSCode和clangd扩展。通过依赖于从Bazel操作图生成的编译数据库，clangd能够为编辑器提供各种智能功能，包括代码补全、跳转到定义等。通过UI或以下命令安装VSCode扩展： ``` code --install-extension ll ```