XLA 2.7 releases don't work properly with the upstream torch 2.7 #8626

hosseinsarshar · 2025-01-26T18:16:33Z

🐛 Bug

I tested the xla 2.7 nightly builds with the upstream torch 2.7 with many date combinations and all resulted in faulty execution - for example:

for this install:

pip install -U --pre torch==2.7.0.dev20250124+cpu --index-url https://download.pytorch.org/whl/nightly/cpu

pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250124-cp310-cp310-linux_x86_64.whl' \
  -f https://storage.googleapis.com/libtpu-releases/index.html \
  -f https://storage.googleapis.com/libtpu-wheels/index.html

pip install -U torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html

I get this behaviour:

$python
Python 3.10.16 (main, Dec 11 2024, 16:24:50) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch_xla
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/hosseins/miniconda3/envs/test-new-nightly/lib/python3.10/site-packages/torch_xla/__init__.py", line 20, in <module>
    import _XLAC
ImportError: /home/hosseins/miniconda3/envs/test-new-nightly/lib/python3.10/site-packages/_XLAC.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN5torch4lazy13MetricFnValueEd

But once, I downgrade the upstream torch to 2.6.0 dev20241216 (it's the highest that worked) - xla works as expected:

pip install -U --pre torch==2.6.0.dev20241216+cpu --index-url https://download.pytorch.org/whl/nightly/cpu

pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250119-cp310-cp310-linux_x86_64.whl' \
  -f https://storage.googleapis.com/libtpu-releases/index.html \
  -f https://storage.googleapis.com/libtpu-wheels/index.html


pip install -U torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html

I get it fixed:

$ python
Python 3.10.16 (main, Dec 11 2024, 16:24:50) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import xla
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'xla'
>>> import torch_xla
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
>>> torch_xla.devices()
[device(type='xla', index=0), device(type='xla', index=1), device(type='xla', index=2), device(type='xla', index=3), device(type='xla', index=4), device(type='xla', index=5), device(type='xla', index=6), device(type='xla', index=7)]

The text was updated successfully, but these errors were encountered:

bhavya01 · 2025-01-27T03:33:27Z

We recently updated the readme with newer instructions for nightly which should fix this issue.

pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl' \
  -f https://storage.googleapis.com/libtpu-releases/index.html \
  -f https://storage.googleapis.com/libtpu-wheels/index.html

# Optional: if you're using custom kernels, install pallas dependencies
pip install 'torch_xla[pallas]' \
  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html

Upstream pytorch changed some environment variables while building their wheels which caused this issue.

hosseinsarshar · 2025-01-27T14:59:02Z

Thanks @bhavya01 - I tried the new instructions as well - but the issue still remains - here is a simple test:

$ conda create -n new-xla python=3.10
Retrieving notices: ...working... done
Channels:
 - defaults
Platform: linux-64
Collecting package metadata (repodata.json): done
Solving environment: done

## Package Plan ##

  environment location: /home/hosseins/miniconda3/envs/new-xla

  added / updated specs:
    - python=3.10


The following NEW packages will be INSTALLED:

  _libgcc_mutex      pkgs/main/linux-64::_libgcc_mutex-0.1-main 
  _openmp_mutex      pkgs/main/linux-64::_openmp_mutex-5.1-1_gnu 
  bzip2              pkgs/main/linux-64::bzip2-1.0.8-h5eee18b_6 
  ca-certificates    pkgs/main/linux-64::ca-certificates-2024.12.31-h06a4308_0 
  ld_impl_linux-64   pkgs/main/linux-64::ld_impl_linux-64-2.40-h12ee557_0 
  libffi             pkgs/main/linux-64::libffi-3.4.4-h6a678d5_1 
  libgcc-ng          pkgs/main/linux-64::libgcc-ng-11.2.0-h1234567_1 
  libgomp            pkgs/main/linux-64::libgomp-11.2.0-h1234567_1 
  libstdcxx-ng       pkgs/main/linux-64::libstdcxx-ng-11.2.0-h1234567_1 
  libuuid            pkgs/main/linux-64::libuuid-1.41.5-h5eee18b_0 
  ncurses            pkgs/main/linux-64::ncurses-6.4-h6a678d5_0 
  openssl            pkgs/main/linux-64::openssl-3.0.15-h5eee18b_0 
  pip                pkgs/main/linux-64::pip-24.2-py310h06a4308_0 
  python             pkgs/main/linux-64::python-3.10.16-he870216_1 
  readline           pkgs/main/linux-64::readline-8.2-h5eee18b_0 
  setuptools         pkgs/main/linux-64::setuptools-75.1.0-py310h06a4308_0 
  sqlite             pkgs/main/linux-64::sqlite-3.45.3-h5eee18b_0 
  tk                 pkgs/main/linux-64::tk-8.6.14-h39e8969_0 
  tzdata             pkgs/main/noarch::tzdata-2025a-h04d1e81_0 
  wheel              pkgs/main/linux-64::wheel-0.44.0-py310h06a4308_0 
  xz                 pkgs/main/linux-64::xz-5.4.6-h5eee18b_1 
  zlib               pkgs/main/linux-64::zlib-1.2.13-h5eee18b_1 


Proceed ([y]/n)? y


Downloading and Extracting Packages:

Preparing transaction: done
Verifying transaction: done
Executing transaction: done
#
# To activate this environment, use
#
#     $ conda activate new-xla
#
# To deactivate an active environment, use
#
#     $ conda deactivate

$ conda activate new-xla
(new-xla) $ pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
Looking in indexes: https://download.pytorch.org/whl/nightly/cpu
Collecting torch
  Downloading https://download.pytorch.org/whl/nightly/cpu/torch-2.7.0.dev20250127%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl.metadata (26 kB)
Collecting torchvision
  Downloading https://download.pytorch.org/whl/nightly/cpu/torchvision-0.22.0.dev20250127%2Bcpu-cp310-cp310-linux_x86_64.whl.metadata (6.2 kB)
Collecting filelock (from torch)
  Using cached https://download.pytorch.org/whl/nightly/filelock-3.16.1-py3-none-any.whl (16 kB)
Collecting typing-extensions>=4.10.0 (from torch)
  Using cached https://download.pytorch.org/whl/nightly/typing_extensions-4.12.2-py3-none-any.whl (37 kB)
Collecting sympy==1.13.1 (from torch)
  Using cached https://download.pytorch.org/whl/nightly/sympy-1.13.1-py3-none-any.whl (6.2 MB)
Collecting networkx (from torch)
  Using cached https://download.pytorch.org/whl/nightly/networkx-3.4.2-py3-none-any.whl (1.7 MB)
Collecting jinja2 (from torch)
  Using cached https://download.pytorch.org/whl/nightly/jinja2-3.1.4-py3-none-any.whl (133 kB)
Collecting fsspec (from torch)
  Using cached https://download.pytorch.org/whl/nightly/fsspec-2024.10.0-py3-none-any.whl (179 kB)
Collecting mpmath<1.4,>=1.1.0 (from sympy==1.13.1->torch)
  Using cached https://download.pytorch.org/whl/nightly/mpmath-1.3.0-py3-none-any.whl (536 kB)
Collecting numpy (from torchvision)
  Using cached https://download.pytorch.org/whl/nightly/numpy-2.1.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (16.3 MB)
Collecting pillow!=8.3.*,>=5.3.0 (from torchvision)
  Using cached https://download.pytorch.org/whl/nightly/pillow-11.0.0-cp310-cp310-manylinux_2_28_x86_64.whl (4.4 MB)
Collecting MarkupSafe>=2.0 (from jinja2->torch)
  Using cached https://download.pytorch.org/whl/nightly/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (25 kB)
Downloading https://download.pytorch.org/whl/nightly/cpu/torch-2.7.0.dev20250127%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl (175.4 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 175.4/175.4 MB 142.6 MB/s eta 0:00:00
Downloading https://download.pytorch.org/whl/nightly/cpu/torchvision-0.22.0.dev20250127%2Bcpu-cp310-cp310-linux_x86_64.whl (1.9 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.9/1.9 MB 133.2 MB/s eta 0:00:00
Installing collected packages: mpmath, typing-extensions, sympy, pillow, numpy, networkx, MarkupSafe, fsspec, filelock, jinja2, torch, torchvision
Successfully installed MarkupSafe-2.1.5 filelock-3.16.1 fsspec-2024.10.0 jinja2-3.1.4 mpmath-1.3.0 networkx-3.4.2 numpy-2.1.2 pillow-11.0.0 sympy-1.13.1 torch-2.7.0.dev20250127+cpu torchvision-0.22.0.dev20250127+cpu typing-extensions-4.12.2
(new-xla) $ pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl' \
  -f https://storage.googleapis.com/libtpu-releases/index.html \
  -f https://storage.googleapis.com/libtpu-wheels/index.html
Looking in links: https://storage.googleapis.com/libtpu-releases/index.html, https://storage.googleapis.com/libtpu-wheels/index.html
Collecting torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl (from torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Downloading https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl (94.6 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 94.6/94.6 MB 146.6 MB/s eta 0:00:00
Collecting absl-py>=1.0.0 (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached absl_py-2.1.0-py3-none-any.whl.metadata (2.3 kB)
Requirement already satisfied: numpy in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl) (2.1.2)
Collecting pyyaml (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached PyYAML-6.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (2.1 kB)
Collecting requests (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached requests-2.32.3-py3-none-any.whl.metadata (4.6 kB)
Collecting libtpu==0.0.8.dev20250113 (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/libtpu-0.0.8.dev20250113%2Bnightly-py3-none-linux_x86_64.whl (132.5 MB)
Collecting tpu-info (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached tpu_info-0.2.0-py3-none-any.whl.metadata (3.7 kB)
Collecting libtpu-nightly==0.1.dev20241010+nightly.cleanup (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20241010%2Bnightly.cleanup-py3-none-any.whl (1.3 kB)
Collecting charset-normalizer<4,>=2 (from requests->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached charset_normalizer-3.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (35 kB)
Collecting idna<4,>=2.5 (from requests->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached idna-3.10-py3-none-any.whl.metadata (10 kB)
Collecting urllib3<3,>=1.21.1 (from requests->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached urllib3-2.3.0-py3-none-any.whl.metadata (6.5 kB)
Collecting certifi>=2017.4.17 (from requests->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached certifi-2024.12.14-py3-none-any.whl.metadata (2.3 kB)
Collecting grpcio>=1.65.5 (from tpu-info->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached grpcio-1.70.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.9 kB)
Collecting protobuf (from tpu-info->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached protobuf-5.29.3-cp38-abi3-manylinux2014_x86_64.whl.metadata (592 bytes)
Collecting rich (from tpu-info->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached rich-13.9.4-py3-none-any.whl.metadata (18 kB)
Collecting markdown-it-py>=2.2.0 (from rich->tpu-info->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached markdown_it_py-3.0.0-py3-none-any.whl.metadata (6.9 kB)
Collecting pygments<3.0.0,>=2.13.0 (from rich->tpu-info->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached pygments-2.19.1-py3-none-any.whl.metadata (2.5 kB)
Requirement already satisfied: typing-extensions<5.0,>=4.0.0 in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from rich->tpu-info->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl) (4.12.2)
Collecting mdurl~=0.1 (from markdown-it-py>=2.2.0->rich->tpu-info->torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Using cached mdurl-0.1.2-py3-none-any.whl.metadata (1.6 kB)
Using cached absl_py-2.1.0-py3-none-any.whl (133 kB)
Using cached PyYAML-6.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (751 kB)
Using cached requests-2.32.3-py3-none-any.whl (64 kB)
Using cached tpu_info-0.2.0-py3-none-any.whl (14 kB)
Using cached certifi-2024.12.14-py3-none-any.whl (164 kB)
Using cached charset_normalizer-3.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (146 kB)
Using cached grpcio-1.70.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.9 MB)
Using cached idna-3.10-py3-none-any.whl (70 kB)
Using cached urllib3-2.3.0-py3-none-any.whl (128 kB)
Using cached protobuf-5.29.3-cp38-abi3-manylinux2014_x86_64.whl (319 kB)
Using cached rich-13.9.4-py3-none-any.whl (242 kB)
Using cached markdown_it_py-3.0.0-py3-none-any.whl (87 kB)
Using cached pygments-2.19.1-py3-none-any.whl (1.2 MB)
Using cached mdurl-0.1.2-py3-none-any.whl (10.0 kB)
Installing collected packages: libtpu-nightly, libtpu, urllib3, pyyaml, pygments, protobuf, mdurl, idna, grpcio, charset-normalizer, certifi, absl-py, requests, markdown-it-py, torch_xla, rich, tpu-info
Successfully installed absl-py-2.1.0 certifi-2024.12.14 charset-normalizer-3.4.1 grpcio-1.70.0 idna-3.10 libtpu-0.0.8.dev20250113+nightly libtpu-nightly-0.1.dev20241010+nightly.cleanup markdown-it-py-3.0.0 mdurl-0.1.2 protobuf-5.29.3 pygments-2.19.1 pyyaml-6.0.2 requests-2.32.3 rich-13.9.4 torch_xla-2.7.0+git8b24140 tpu-info-0.2.0 urllib3-2.3.0
(new-xla) $ pip install 'torch_xla[pallas]' \
  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
Looking in links: https://storage.googleapis.com/jax-releases/jax_nightly_releases.html, https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
Requirement already satisfied: torch_xla[pallas] in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (2.7.0+git8b24140)
Requirement already satisfied: absl-py>=1.0.0 in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from torch_xla[pallas]) (2.1.0)
Requirement already satisfied: numpy in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from torch_xla[pallas]) (2.1.2)
Requirement already satisfied: pyyaml in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from torch_xla[pallas]) (6.0.2)
Requirement already satisfied: requests in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from torch_xla[pallas]) (2.32.3)
Collecting jaxlib==0.4.39.dev20250113 (from torch_xla[pallas])
  Using cached https://storage.googleapis.com/jax-releases/nightly/nocuda/jaxlib-0.4.39.dev20250113-cp310-cp310-manylinux2014_x86_64.whl (101.5 MB)
Collecting jax==0.4.39.dev20250113 (from torch_xla[pallas])
  Using cached https://storage.googleapis.com/jax-releases/nightly/jax/jax-0.4.39.dev20250113-py3-none-any.whl (2.3 MB)
Collecting ml_dtypes>=0.4.0 (from jax==0.4.39.dev20250113->torch_xla[pallas])
  Using cached ml_dtypes-0.5.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (21 kB)
Collecting opt_einsum (from jax==0.4.39.dev20250113->torch_xla[pallas])
  Using cached opt_einsum-3.4.0-py3-none-any.whl.metadata (6.3 kB)
Collecting scipy>=1.11.1 (from jax==0.4.39.dev20250113->torch_xla[pallas])
  Using cached scipy-1.15.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (61 kB)
Requirement already satisfied: charset-normalizer<4,>=2 in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from requests->torch_xla[pallas]) (3.4.1)
Requirement already satisfied: idna<4,>=2.5 in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from requests->torch_xla[pallas]) (3.10)
Requirement already satisfied: urllib3<3,>=1.21.1 in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from requests->torch_xla[pallas]) (2.3.0)
Requirement already satisfied: certifi>=2017.4.17 in ./miniconda3/envs/new-xla/lib/python3.10/site-packages (from requests->torch_xla[pallas]) (2024.12.14)
Using cached ml_dtypes-0.5.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.7 MB)
Using cached scipy-1.15.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (40.6 MB)
Using cached opt_einsum-3.4.0-py3-none-any.whl (71 kB)
Installing collected packages: scipy, opt_einsum, ml_dtypes, jaxlib, jax
Successfully installed jax-0.4.39.dev20250113 jaxlib-0.4.39.dev20250113 ml_dtypes-0.5.1 opt_einsum-3.4.0 scipy-1.15.1

(new-xla) $ python
Python 3.10.16 (main, Dec 11 2024, 16:24:50) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch_xla
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/hosseins/miniconda3/envs/new-xla/lib/python3.10/site-packages/torch_xla/__init__.py", line 20, in <module>
    import _XLAC
ImportError: /home/hosseins/miniconda3/envs/new-xla/lib/python3.10/site-packages/_XLAC.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN5torch6dynamo8autograd18get_input_metadataERKSt6vectorINS_8autograd4EdgeESaIS4_EE
>>>

miladm · 2025-01-27T15:36:12Z

Ideally, this build combination should work as expected @hosseinsarshar.

I see you were able to unblock this issue by matching an older torch wheel with the latest torch_xla. I am reopening it so we can further investigate further. cc @lsy323 @ysiraichi

hosseinsarshar · 2025-02-03T21:58:09Z

@miladm / @lsy323 wonder if you had a chance to review this issue - I'm asking as the current pinned torch version (2.6.0.dev20241216+cpu) is out of date today. Thanks

ysiraichi · 2025-02-06T13:47:56Z

I couldn't reproduce this issue. I ran the same commands (using micromamba, instead), and I was able to import torch_xla without problems.

ysiraichi · 2025-02-06T13:55:57Z

~
➜ micromamba create -n test python=3.10
conda-forge/noarch                                  18.9MB @  19.5MB/s  1.0s
conda-forge/linux-64                                41.7MB @  15.0MB/s  2.8s


Transaction
  ...
  Updating specs:

   - python=3.10


  Package               Version  Build               Channel           Size
─────────────────────────────────────────────────────────────────────────────
  Install:
─────────────────────────────────────────────────────────────────────────────

  + _libgcc_mutex           0.1  conda_forge         conda-forge     Cached
  + _openmp_mutex           4.5  2_gnu               conda-forge     Cached
  + bzip2                 1.0.8  h4bc722e_7          conda-forge     Cached
  + ca-certificates   2025.1.31  hbcca054_0          conda-forge     Cached
  + ld_impl_linux-64       2.43  h712a8e2_2          conda-forge     Cached
  + libffi                3.4.2  h7f98852_5          conda-forge     Cached
  + libgcc               14.2.0  h77fa898_1          conda-forge     Cached
  + libgcc-ng            14.2.0  h69a702a_1          conda-forge     Cached
  + libgomp              14.2.0  h77fa898_1          conda-forge     Cached
  + liblzma               5.6.4  hb9d3cd8_0          conda-forge     Cached
  + libnsl                2.0.1  hd590300_0          conda-forge     Cached
  + libsqlite            3.48.0  hee588c1_1          conda-forge     Cached
  + libuuid              2.38.1  h0b41bf4_0          conda-forge     Cached
  + libxcrypt            4.4.36  hd590300_1          conda-forge     Cached
  + libzlib               1.3.1  hb9d3cd8_2          conda-forge     Cached
  + ncurses                 6.5  h2d0b736_3          conda-forge     Cached
  + openssl               3.4.0  h7b32b05_1          conda-forge     Cached
  + pip                    25.0  pyh8b19718_0        conda-forge     Cached
  + python              3.10.16  he725a3c_1_cpython  conda-forge     Cached
  + readline                8.2  h8228510_1          conda-forge     Cached
  + setuptools           75.8.0  pyhff2d567_0        conda-forge     Cached
  + tk                   8.6.13  noxft_h4845f30_101  conda-forge     Cached
  + tzdata                2025a  h78e105d_0          conda-forge     Cached
  + wheel                0.45.1  pyhd8ed1ab_1        conda-forge     Cached

  Summary:

  Install: 24 packages

  Total download: 0 B

─────────────────────────────────────────────────────────────────────────────


Confirm changes: [Y/n]
...

~
➜ micromamba activate test

~ via 🅒 test
➜ pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
Looking in indexes: https://download.pytorch.org/whl/nightly/cpu
Collecting torch
  Downloading https://download.pytorch.org/whl/nightly/cpu/torch-2.7.0.dev20250206%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl.metadata (26 kB)
Collecting torchvision
  Downloading https://download.pytorch.org/whl/nightly/cpu/torchvision-0.22.0.dev20250206%2Bcpu-cp310-cp310-linux_x86_64.whl.metadata (6.2 kB)
...

~ via 🅒 test
➜ pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl' \
  -f https://storage.googleapis.com/libtpu-releases/index.html \
  -f https://storage.googleapis.com/libtpu-wheels/index.html
Looking in links: https://storage.googleapis.com/libtpu-releases/index.html, https://storage.googleapis.com/libtpu-wheels/index.html
Collecting torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl (from torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Downloading https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl (95.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 95.0/95.0 MB 19.7 MB/s eta 0:00:00
Collecting libtpu==0.0.9.dev20250131 (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Downloading https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/libtpu-0.0.9.dev20250131%2Bnightly-py3-none-linux_x86_64.whl (133.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 133.1/133.1 MB 28.7 MB/s eta 0:00:00
Collecting tpu-info (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Downloading tpu_info-0.2.0-py3-none-any.whl.metadata (3.7 kB)
Collecting libtpu-nightly==0.1.dev20241010+nightly.cleanup (from torch_xla@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl->torch_xla[tpu]@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev+cxx11-cp310-cp310-linux_x86_64.whl)
  Downloading https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20241010%2Bnightly.cleanup-py3-none-any.whl (1.3 kB)
...

~ via 🅒 test
➜ pip install 'torch_xla[pallas]' \
  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
Looking in links: https://storage.googleapis.com/jax-releases/jax_nightly_releases.html, https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
Collecting jaxlib==0.5.1.dev20250131 (from torch_xla[pallas])
  Downloading https://storage.googleapis.com/jax-releases/nightly/nocuda/jaxlib-0.5.1.dev20250131-cp310-cp310-manylinux2014_x86_64.whl (103.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 103.0/103.0 MB 22.5 MB/s eta 0:00:00
Collecting jax==0.5.1.dev20250131 (from torch_xla[pallas])
  Downloading https://storage.googleapis.com/jax-releases/nightly/jax/jax-0.5.1.dev20250131-py3-none-any.whl (2.3 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.3/2.3 MB 5.0 MB/s eta 0:00:00
Collecting ml_dtypes>=0.4.0 (from jax==0.5.1.dev20250131->torch_xla[pallas])
  Downloading ml_dtypes-0.5.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (21 kB)
...

~ via 🅒 test took 15s509ms
➜ python
Python 3.10.16 | packaged by conda-forge | (main, Dec  5 2024, 14:16:10) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch_xla
WARNING:root:Defaulting to PJRT_DEVICE=CPU

hosseinsarshar · 2025-02-06T16:44:08Z

thanks @ysiraichi for testing this -

Checking your logs - I see WARNING:root:Defaulting to PJRT_DEVICE=CPU which shows that it can't find the TPU device by default which shouldn't be the case. I also mentioned it in my first comment - if you pass PJRT_DEVICE=TPU, it most probably break - if it doesn't, I expect you can't get the list of devices: please try to list them by torch_xla.devices().

hosseinsarshar changed the title ~~XLA 2.7 releases don't work properly with upstream torch 2.7~~ XLA 2.7 releases don't work properly with the upstream torch 2.7 Jan 26, 2025

hosseinsarshar mentioned this issue Jan 26, 2025

[CI/Build] Fixed the xla nightly issue report in #12451 vllm-project/vllm#12453

Merged

bhavya01 closed this as completed Jan 27, 2025

miladm assigned ysiraichi Jan 27, 2025

miladm reopened this Jan 27, 2025

ysiraichi added needs reproduction xla:tpu TPU specific issues and PRs labels Feb 11, 2025

iwknow mentioned this issue Mar 10, 2025

undefined symbol from _XLAC #8813

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

XLA 2.7 releases don't work properly with the upstream torch 2.7 #8626

XLA 2.7 releases don't work properly with the upstream torch 2.7 #8626

hosseinsarshar commented Jan 26, 2025 •

edited

Loading

bhavya01 commented Jan 27, 2025

hosseinsarshar commented Jan 27, 2025

miladm commented Jan 27, 2025 •

edited

Loading

hosseinsarshar commented Feb 3, 2025 •

edited

Loading

ysiraichi commented Feb 6, 2025

ysiraichi commented Feb 6, 2025

hosseinsarshar commented Feb 6, 2025

XLA 2.7 releases don't work properly with the upstream torch 2.7 #8626

XLA 2.7 releases don't work properly with the upstream torch 2.7 #8626

Comments

hosseinsarshar commented Jan 26, 2025 • edited Loading

🐛 Bug

bhavya01 commented Jan 27, 2025

hosseinsarshar commented Jan 27, 2025

miladm commented Jan 27, 2025 • edited Loading

hosseinsarshar commented Feb 3, 2025 • edited Loading

ysiraichi commented Feb 6, 2025

ysiraichi commented Feb 6, 2025

hosseinsarshar commented Feb 6, 2025

hosseinsarshar commented Jan 26, 2025 •

edited

Loading

miladm commented Jan 27, 2025 •

edited

Loading

hosseinsarshar commented Feb 3, 2025 •

edited

Loading