`undefined symbol: _ZN3c104cuda10set_deviceEa` when trying to use torch_xla 2.8

[jrfonseca]

🐛 Bug

I've been using torch_xla 2.7 without issues, but when I try to use torch_xla 2.8 I run into undefined symbol: _ZN3c104cuda10set_deviceEa errors.

The corresponding c++filt demangled symbol name is c10::cuda::set_device(signed char)

To Reproduce

Steps to reproduce the behavior:

1. Create a repro.sh containing:

#!/bin/sh

set -eu

# Using https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa
for pyver in 3.11 3.12 3.13
do
        python$pyver -m venv .venv-py$pyver
        . .venv-py$pyver/bin/activate
        test -d "$VIRTUAL_ENV"
        uv pip install -q torch==2.8.0 'torch_xla[tpu]==2.8.0'
        pyver_major=$(python -c 'import sys; print(sys.version_info[0])')
        pyver_minor=$(python -c 'import sys; print(sys.version_info[1])')
        LD_LIBRARY_PATH=$VIRTUAL_ENV/lib/python${pyver_major}.${pyver_minor}/site-packages/torch/lib ldd -r $VIRTUAL_ENV/lib/python${pyver_major}.${pyver_minor}/site-packages/_XLAC.cpython-${pyver_major}${pyver_minor}-x86_64-linux-gnu.so
        python -c 'import torch_xla' || true
        deactivate
done

2. Run it as

$ ./repro.sh > /dev/null
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File ".../.venv-py3.11/lib/python3.11/site-packages/torch_xla/__init__.py", line 20, in <module>
    import _XLAC
ImportError: .../.venv-py3.11/lib/python3.11/site-packages/_XLAC.cpython-311-x86_64-linux-gnu.so: undefined symbol: _ZN3c104cuda10set_deviceEa
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File ".../.venv-py3.12/lib/python3.12/site-packages/torch_xla/__init__.py", line 20, in <module>
    import _XLAC
ImportError: .../.venv-py3.12/lib/python3.12/site-packages/_XLAC.cpython-312-x86_64-linux-gnu.so: undefined symbol: _ZN3c104cuda10set_deviceEa
Traceback (most recent call last):
  File "<string>", line 1, in <module>
    import torch_xla
  File ".../.venv-py3.13/lib/python3.13/site-packages/torch_xla/__init__.py", line 20, in <module>
    import _XLAC
ImportError: .../.venv-py3.13/lib/python3.13/site-packages/_XLAC.cpython-313-x86_64-linux-gnu.so: undefined symbol: _ZN3c104cuda10set_deviceEa

Expected behavior

torch_xla successfully imported without missing symbols errors

Environment

• Ubuntu 24.04 (includes Python 3.12)
• Additional Python versions from https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa
• Reproducible on XLA backend [CPU/TPU]: both?
• torch_xla version: 2.8

Additional context

This happens no matter what Python version I use.

I also tried to use the torch_xla packages from https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/... without success.

It seems that torch_xla packages were built against a pytorch binaries significantly different from the ones available from PyPi?

[jrfonseca]

I noticed that libc10_cuda.so exports a very similar symbol:

$ nm --dynamic --defined-only .venv-py3.11/lib/python3.11/site-packages/torch/lib/libc10_cuda.so   | grep _ZN3c104cuda10set_deviceEa
00000000000535b0 T _ZN3c104cuda10set_deviceEab
$ echo _ZN3c104cuda10set_deviceEab | c++filt
c10::cuda::set_device(signed char, bool)
$ echo _ZN3c104cuda10set_deviceEa | c++filt
c10::cuda::set_device(signed char)

[jrfonseca]

A few more observations:

1. c10::cuda::set_device() prototype was modified upstream on pytorch/pytorch@4a26bb8 . I don't understand how torch_xla was built agains pytorch 2.8.0 yet didn't pick up the right prototype.

2. torch_xla appears to include a stub on https://github.com/pytorch/xla/blob/v2.8.0/torch_xla/csrc/aten_cuda_functions.cpp#L45 and

   xla/torch_xla/csrc/aten_cuda_functions.h

   Line 27 in d517649

   void set_device(c10::DeviceIndex);

   that supposedly should be used when used with torch[cpu], but that stub is using the wrong prototype. I think this is the problem.

3. @ysiraichi's e9a1c5f removed this code, so this is only a problem in torch_xla v2.8.x

[jrfonseca]

I believe something like below (untested) should fix the issue.

diff --git a/torch_xla/csrc/aten_cuda_functions.cpp b/torch_xla/csrc/aten_cuda_functions.cpp
index 45a5a4a35..f4490734e 100644
--- a/torch_xla/csrc/aten_cuda_functions.cpp
+++ b/torch_xla/csrc/aten_cuda_functions.cpp
@@ -42,7 +42,7 @@ DeviceIndex device_count() noexcept { return 0; }

 c10::DeviceIndex current_device() { fail("c10::cuda::current_device()"); }

-void set_device(c10::DeviceIndex) { fail("c10::cuda::set_device()"); }
+void set_device(c10::DeviceIndex, const bool force) { fail("c10::cuda::set_device()"); }

 void device_synchronize() { fail("c10::cuda::device_synchronize()"); }

diff --git a/torch_xla/csrc/aten_cuda_functions.h b/torch_xla/csrc/aten_cuda_functions.h
index 50cb3de1a..4813788cc 100644
--- a/torch_xla/csrc/aten_cuda_functions.h
+++ b/torch_xla/csrc/aten_cuda_functions.h
@@ -24,7 +24,7 @@ DeviceIndex device_count() noexcept;

 c10::DeviceIndex current_device();

-void set_device(c10::DeviceIndex);
+void set_device(c10::DeviceIndex, const bool force = false);

 void device_synchronize();

[jrfonseca]

In theory if one followed https://pytorch.org/ instructions and installed CPU build of torch this issue would not happen.

Perhaps that's why this issue went unnoticed for so long.

Unfortunately, in practice https://pytorch.org/ instructions are not working: when I run pip install torch I get a torch build with CUDA support. Perhaps this is a recent development / regression in upstream pytorch, which made this torch_xla bug surface.

Unfortunately I still found no workaround to this, other than rebuilding pytorch and/or torch_xla from source, which I really really was hoping I could avoid...

[jrfonseca]

I finally figure out how to get a CPU torch build installed:

$ uv pip install 'torch==2.8.0+cpu' 'torch_xla==2.8.0' --extra-index-url https://download.pytorch.org/whl/cpu
[...]
$ python -c 'import torch.cuda; print(torch.cuda.is_available()); import torch_xla'
False
WARNING:root:Defaulting to PJRT_DEVICE=CPU

So, if you don't want to make another torch_xla 2.8.x release, perhaps you can update the documentation so others can avoid falling into this deep rabbit hole I felt into!

[ysiraichi]

Thank you for filing this issue.
This sure was quite the rabbit hole. Again, thank you for documenting it, here.

Issue Explanation: the problem was that aten_cuda_functions.{h,cpp} were created as a fallback implementation to CUDA functions, whenever the underlying PyTorch didn't support CUDA. For PyTorch installations that supported CUDA, we would use the definitions provided by PyTorch itself. However, when upstream changed the signature of those functions, the declarations at aten_cuda_functions.h were different from the ones provided by PyTorch. In other words, when we tried to call c10::cuda::set_device(), those "new" CUDA functions were nowhere to be found. Thus, the error.

Why does it work with PyTorch CPU? aten_cuda_functions.cpp is compiled into another Python extension that we conditionally load only for PyTorch CPU. So, when we load that extension, everything works. When we don't (PyTorch with CUDA support), we get that error.

I think we should probably fix that with your proposed solution.
Could you submit a PR?

[jrfonseca]

Thanks @ysiraichi for looking into this and the summary.

Could you submit a PR?

I'd appreciate if you could independently fix the best way you see fit. Although I believe that the changes proposed in #9656 (comment) are too trivial to be copyrightable, I'm not quite ready yet to formally contribute code to this project. Sorry for inconvenience this may cause.