From f06b3e557be4e91a440c4bb060594ca64b130e7c Mon Sep 17 00:00:00 2001 From: ejmartinezm Date: Wed, 23 Sep 2026 14:36:51 -0700 Subject: [PATCH 1/6] Add TorchTPU entry --- _additional_platforms/torch_tpu.json | 7 +++ .../additional_platforms/torch_tpu.md | 58 +++++++++++++++++++ 2 files changed, 65 insertions(+) create mode 100644 _additional_platforms/torch_tpu.json create mode 100644 _get_started/additional_platforms/torch_tpu.md diff --git a/_additional_platforms/torch_tpu.json b/_additional_platforms/torch_tpu.json new file mode 100644 index 000000000000..de55b5fbdf3c --- /dev/null +++ b/_additional_platforms/torch_tpu.json @@ -0,0 +1,7 @@ +{ + "name": "TorchTPU", + "support_channel": "https://github.com/google-pytorch/torch_tpu/issues", + "stable": { + "linux": "pip3 install torch --index-url https://download.pytorch.org/whl/cpu && pip3 install torch-tpu" + } +} \ No newline at end of file diff --git a/_get_started/additional_platforms/torch_tpu.md b/_get_started/additional_platforms/torch_tpu.md new file mode 100644 index 000000000000..63c4361a1d61 --- /dev/null +++ b/_get_started/additional_platforms/torch_tpu.md @@ -0,0 +1,58 @@ +# Installing on Google Cloud TPU (TorchTPU) + +**TorchTPU (`torch_tpu`)** is a native PyTorch backend built for Google Cloud Tensor Processing Units (TPUs). It enables Google Cloud TPUs to run PyTorch workloads natively (`PrivateUse1` dispatch key `"tpu"`), supporting eager execution (`DeferAndFuse`), `torch.compile()`, distributed training (`torch.distributed`, `DTensor`, `FSDP2`), and custom kernels (`Pallas`, `Helion`). + +## Prerequisites + +### Hardware Requirements + +* A provisioned Google Cloud TPU VM (TPU v5e, v5p, v6e, or v7x) on Google Compute Engine (GCE) or Google Kubernetes Engine (GKE), or a TPU runtime in Google Colab. + +### Software Requirements + +* Linux (Ubuntu 22.04+ recommended) +* Python >= 3.10 +* `libtpu` runtime library (automatically installed with `torch-tpu`) +* **Recommended Memory Allocator**: For optimal runtime performance, configure `TCMalloc` (`LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4`) prior to launching Python workloads, following the [PyTorch Performance Tuning Guide](https://docs.pytorch.org/tutorials/recipes/recipes/tuning_guide.html#switch-memory-allocator). + +## Installation + +### pip + +```bash +pip3 install torch --index-url https://download.pytorch.org/whl/cpu && pip3 install torch-tpu +``` + +Use the `pip` package manager to install the CPU build of PyTorch alongside `torch-tpu`. Select your preferred options in the selector above to get the installation command. + +## Verification + +To ensure that PyTorch was installed correctly with TorchTPU support, run the following code: + +```python +import torch +import torch_tpu + +x = torch.randn(2, 2, device="tpu") +y = torch.randn(2, 2, device="tpu") +z = x.mm(y) + +print(z) +``` + +The following, or a similar output, indicates successful installation: + +```bash +tensor([[-0.4218, 0.8912], + [ 0.1534, -1.1045]], device='tpu:0') +``` + +## Documentation + +For more information, please visit: + +* [TorchTPU Official Documentation & User Guide](https://google-pytorch.github.io/torch_tpu/) +* [Supported vs. Unsupported Feature Matrix](https://github.com/google-pytorch/torch_tpu/blob/main/docs/features_matrix.md) +* [Google Cloud TPU Documentation](https://cloud.google.com/tpu/docs) +* [GitHub Repository](https://github.com/google-pytorch/torch_tpu) +* [PyPI](https://pypi.org/project/torch-tpu/) \ No newline at end of file From 6503a32d071a8e3cfd8c125c344f64dda04cedd1 Mon Sep 17 00:00:00 2001 From: ejmartinezm Date: Wed, 23 Sep 2026 14:47:19 -0700 Subject: [PATCH 2/6] Update TorchTPU installation documentation Fixed formatting in the TorchTPU documentation and corrected the device assignment example. --- _get_started/additional_platforms/torch_tpu.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/_get_started/additional_platforms/torch_tpu.md b/_get_started/additional_platforms/torch_tpu.md index 63c4361a1d61..2b7153b50c36 100644 --- a/_get_started/additional_platforms/torch_tpu.md +++ b/_get_started/additional_platforms/torch_tpu.md @@ -1,6 +1,6 @@ # Installing on Google Cloud TPU (TorchTPU) -**TorchTPU (`torch_tpu`)** is a native PyTorch backend built for Google Cloud Tensor Processing Units (TPUs). It enables Google Cloud TPUs to run PyTorch workloads natively (`PrivateUse1` dispatch key `"tpu"`), supporting eager execution (`DeferAndFuse`), `torch.compile()`, distributed training (`torch.distributed`, `DTensor`, `FSDP2`), and custom kernels (`Pallas`, `Helion`). +**TorchTPU (`torch_tpu`)** is a native PyTorch backend built for Google Cloud Tensor Processing Units (TPUs). It enables Google Cloud TPUs to run PyTorch workloads natively (`PrivateUse1` dispatch key `"tpu"`), supporting eager execution `torch.compile()`, distributed training (`torch.distributed`, `DTensor`, `FSDP2`), and custom kernels (`Pallas`, `Helion`). ## Prerequisites @@ -31,7 +31,7 @@ To ensure that PyTorch was installed correctly with TorchTPU support, run the fo ```python import torch -import torch_tpu +device =torch.device("tpu") x = torch.randn(2, 2, device="tpu") y = torch.randn(2, 2, device="tpu") @@ -55,4 +55,4 @@ For more information, please visit: * [Supported vs. Unsupported Feature Matrix](https://github.com/google-pytorch/torch_tpu/blob/main/docs/features_matrix.md) * [Google Cloud TPU Documentation](https://cloud.google.com/tpu/docs) * [GitHub Repository](https://github.com/google-pytorch/torch_tpu) -* [PyPI](https://pypi.org/project/torch-tpu/) \ No newline at end of file +* [PyPI](https://pypi.org/project/torch-tpu/) From 235c29ae1969c54c715ed4539aca82f47836d2e8 Mon Sep 17 00:00:00 2001 From: ejmartinezm Date: Fri, 25 Sep 2026 16:01:20 -0700 Subject: [PATCH 3/6] Update _get_started/additional_platforms/torch_tpu.md Co-authored-by: LeeHowes --- _get_started/additional_platforms/torch_tpu.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/_get_started/additional_platforms/torch_tpu.md b/_get_started/additional_platforms/torch_tpu.md index 2b7153b50c36..99a9a59a4073 100644 --- a/_get_started/additional_platforms/torch_tpu.md +++ b/_get_started/additional_platforms/torch_tpu.md @@ -1,6 +1,6 @@ # Installing on Google Cloud TPU (TorchTPU) -**TorchTPU (`torch_tpu`)** is a native PyTorch backend built for Google Cloud Tensor Processing Units (TPUs). It enables Google Cloud TPUs to run PyTorch workloads natively (`PrivateUse1` dispatch key `"tpu"`), supporting eager execution `torch.compile()`, distributed training (`torch.distributed`, `DTensor`, `FSDP2`), and custom kernels (`Pallas`, `Helion`). +**TorchTPU (`torch_tpu`)** is a native PyTorch backend built for Google Cloud Tensor Processing Units (TPUs). It enables Google Cloud TPUs to run PyTorch workloads natively by integrating as a `PrivateUse1` backend using the dispatch key `"tpu"`. TorchTPU supports eager execution, `torch.compile()`, distributed training (`torch.distributed`, `DTensor`, `FSDP2`), and custom kernels (`Pallas`, `Helion`). ## Prerequisites From 35796c3623719dc2feabbd01f23d420a451fc22a Mon Sep 17 00:00:00 2001 From: ejmartinezm Date: Fri, 25 Sep 2026 16:01:42 -0700 Subject: [PATCH 4/6] Update _get_started/additional_platforms/torch_tpu.md Co-authored-by: LeeHowes --- _get_started/additional_platforms/torch_tpu.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/_get_started/additional_platforms/torch_tpu.md b/_get_started/additional_platforms/torch_tpu.md index 99a9a59a4073..2a3525a16b67 100644 --- a/_get_started/additional_platforms/torch_tpu.md +++ b/_get_started/additional_platforms/torch_tpu.md @@ -6,7 +6,7 @@ ### Hardware Requirements -* A provisioned Google Cloud TPU VM (TPU v5e, v5p, v6e, or v7x) on Google Compute Engine (GCE) or Google Kubernetes Engine (GKE), or a TPU runtime in Google Colab. +* A provisioned Google Cloud TPU VM (TPU v5e, v5p, v6e, or v7x) on Google Compute Engine (GCE), access to Google Kubernetes Engine (GKE) cluster, or a TPU runtime in Google Colab. ### Software Requirements From ec9b17a485e5c8bce69e31ddf18d111aa853e881 Mon Sep 17 00:00:00 2001 From: ejmartinezm Date: Fri, 25 Sep 2026 16:03:11 -0700 Subject: [PATCH 5/6] Remove TCMalloc recommendation from torch_tpu.md Removed the recommendation for TCMalloc memory allocator for optimal runtime performance. --- _get_started/additional_platforms/torch_tpu.md | 1 - 1 file changed, 1 deletion(-) diff --git a/_get_started/additional_platforms/torch_tpu.md b/_get_started/additional_platforms/torch_tpu.md index 2a3525a16b67..126437b5d3b7 100644 --- a/_get_started/additional_platforms/torch_tpu.md +++ b/_get_started/additional_platforms/torch_tpu.md @@ -13,7 +13,6 @@ * Linux (Ubuntu 22.04+ recommended) * Python >= 3.10 * `libtpu` runtime library (automatically installed with `torch-tpu`) -* **Recommended Memory Allocator**: For optimal runtime performance, configure `TCMalloc` (`LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4`) prior to launching Python workloads, following the [PyTorch Performance Tuning Guide](https://docs.pytorch.org/tutorials/recipes/recipes/tuning_guide.html#switch-memory-allocator). ## Installation From 5297ccc021a416e449e17cee4ff58896f539d145 Mon Sep 17 00:00:00 2001 From: ejmartinezm <18467846+ejmartinezm@users.noreply.github.com> Date: Tue, 29 Sep 2026 16:47:29 -0700 Subject: [PATCH 6/6] Update _get_started/additional_platforms/torch_tpu.md Co-authored-by: albanD --- _get_started/additional_platforms/torch_tpu.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/_get_started/additional_platforms/torch_tpu.md b/_get_started/additional_platforms/torch_tpu.md index 126437b5d3b7..afc23e782b49 100644 --- a/_get_started/additional_platforms/torch_tpu.md +++ b/_get_started/additional_platforms/torch_tpu.md @@ -30,7 +30,7 @@ To ensure that PyTorch was installed correctly with TorchTPU support, run the fo ```python import torch -device =torch.device("tpu") +device = torch.device("tpu") x = torch.randn(2, 2, device="tpu") y = torch.randn(2, 2, device="tpu")