@@ -26,7 +26,7 @@ started:
26
26
To install PyTorch/XLA stable build in a new TPU VM:
27
27
28
28
``` sh
29
- pip install torch~=2.5 .0 ' torch_xla[tpu]~=2.5 .0' \
29
+ pip install torch~=2.6 .0 ' torch_xla[tpu]~=2.6 .0' \
30
30
-f https://storage.googleapis.com/libtpu-releases/index.html \
31
31
-f https://storage.googleapis.com/libtpu-wheels/index.html
32
32
@@ -50,6 +50,36 @@ pip install 'torch_xla[pallas]' \
50
50
-f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
51
51
```
52
52
53
+ ### C++11 ABI builds
54
+
55
+ Starting from Pytorch/XLA 2.6, we'll provide wheels and docker images built with
56
+ two C++ ABI flavors: C++11 and pre-C++11. Pre-C++11 is the default to align with
57
+ PyTorch upstream, but C++11 ABI wheels and docker images have better lazy tensor
58
+ tracing performance.
59
+
60
+ To install C++11 ABI flavored 2.6 wheels:
61
+
62
+ ``` sh
63
+ pip install torch==2.6.0+cpu.cxx11.abi torch_xla[tpu]==2.6.0+cxx11 \
64
+ -f https://storage.googleapis.com/libtpu-releases/index.html \
65
+ -f https://storage.googleapis.com/libtpu-wheels/index.html \
66
+ -f https://download.pytorch.org/whl/torch
67
+ ```
68
+
69
+ To access C++11 ABI flavored docker image:
70
+
71
+ ```
72
+ us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11
73
+ ```
74
+
75
+ If your model is tracing bound (e.g. you see that the host CPU is busy tracing
76
+ the model while TPUs are idle), switching to the C++11 ABI wheels/docker images
77
+ can improve performance. Mixtral 8x7B benchmarking results on v5p-256, global
78
+ batch size 1024:
79
+
80
+ - Pre-C++11 ABI MFU: 33%
81
+ - C++ ABI MFU: 39%
82
+
53
83
### GPU Plugin
54
84
55
85
PyTorch/XLA now provides GPU support through a plugin package similar to ` libtpu ` :
@@ -58,6 +88,13 @@ PyTorch/XLA now provides GPU support through a plugin package similar to `libtpu
58
88
pip install torch~=2.5.0 torch_xla~=2.5.0 https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla_cuda_plugin-2.5.0-py3-none-any.whl
59
89
```
60
90
91
+ The newest stable version where PyTorch/XLA: GPU wheel is available is ` torch_xla `
92
+ 2.5. We do not offer a PyTorch/XLA: GPU wheel in the PyTorch/XLA 2.6 release. We
93
+ understand this is important and plan to [ reinstate GPU support] ( https://github.com/pytorch/xla/issues/8577 ) by the 2.7 release.
94
+ PyTorch/XLA remains an open-source project and we welcome contributions from the
95
+ community to help maintain and improve the project. To contribute, please start
96
+ with the [ contributors guide] ( https://github.com/pytorch/xla/blob/master/CONTRIBUTING.md ) .
97
+
61
98
## Getting Started
62
99
63
100
To update your existing training loop, make the following changes:
@@ -203,6 +240,7 @@ The torch wheel version `2.7.0.dev20250124+cpu` can be found at https://download
203
240
204
241
| Version | Cloud TPU VMs Wheel |
205
242
| ---------| -------------------|
243
+ | 2.5 (Python 3.10) | ` https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0-cp310-cp310-manylinux_2_28_x86_64.whl ` |
206
244
| 2.4 (Python 3.10) | ` https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.4.0-cp310-cp310-manylinux_2_28_x86_64.whl ` |
207
245
| 2.3 (Python 3.10) | ` https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.3.0-cp310-cp310-manylinux_2_28_x86_64.whl ` |
208
246
| 2.2 (Python 3.10) | ` https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.2.0-cp310-cp310-manylinux_2_28_x86_64.whl ` |
@@ -236,6 +274,7 @@ The torch wheel version `2.7.0.dev20250124+cpu` can be found at https://download
236
274
237
275
| Version | Cloud TPU VMs Docker |
238
276
| --- | ----------- |
277
+ | 2.6 | ` us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm ` |
239
278
| 2.5 | ` us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.5.0_3.10_tpuvm ` |
240
279
| 2.4 | ` us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm ` |
241
280
| 2.3 | ` us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.3.0_3.10_tpuvm ` |
0 commit comments