Support XLA_GPU and CUDA cowork by default

[kevint324]

🚀 Feature

The feature request is pretty much the same as this issue Running regular PyTorch with PT-XLA CUDA Container.

Motivation

It is a little bit counter intuititve that XLA CUDA and CUDA cannot coexist.
Usually XLA_GPU users and CUDA users are highly overlapped.
And a very common scenario is to compare performance between the native CUDA mode and XLA mode.
Due to current design limitation we cannot do it within the same docker.

Support XLA_GPU/CUDA cowork will greatly improve PyTorch/XLA's usability and user experience.

Pitch

1. Turn on CUDA support in PyTorch/XLA docker by default
2. Add certain tests to ensure XLA_GPU and CUDA could work together.

Alternatives

Additional context

Same confusion raised in this post

We disabled the pytorch wheel with USE_CUDA=0 since there was a old bug that we can’t have both pytorch and pytorch/xla both trying to grab the CUDA device. That issue might or might not still be the case.

Same confusion 2

#3452

[JackCaoG]

Enabling the CUDA build for pytorch should be easy. For moving tensor between CUDA and XLA:GPU I believe @cicirori has some ideas in #4692

[qipengh]

@JackCaoG I also do some benchmark test of performance with XLA_GPU and Torch_CUDA， like #4699 。I see torch_xla does some update code of PJRT after release 1.13, so I want to try new features with r2.0 or master 。And I install torch_xla and torch within GPU docker of CUDA 11.7 as follow:

$ pip install https://download.pytorch.org/whl/nightly/torch/torch-2.0.0.dev20230226+cu117-cp38-cp38-linux_x86_64.whl
$ pip install https://storage.googleapis.com/tpu-pytorch/wheels/cuda/117/torch_xla-nightly+20230228-cp38-cp38-linux_x86_64.whl

But I got error with _XLAC.cpython-38-x86_64-linux-gnu.so: undefined symbol: _ZNK5torch8autograd4Node4nameB5cxx11Ev

(pytorch) root@PowerEdge-C4140:/projs/framework/huangqipeng/pytorch/aa/xla/xla/test# PJRT_DEVICE=GPU GPU_NUM_DEVICES=1 python -c "import torch_xla;torch_xla._XLAC._xla_get_default_device()"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch_xla/__init__.py", line 134, in <module>
    import _XLAC
ImportError: /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/_XLAC.cpython-38-x86_64-linux-gnu.so: undefined symbol: _ZNK5torch8autograd4Node4nameB5cxx11Ev
(pytorch) root@PowerEdge-C4140:/projs/framework/huangqipeng/pytorch/aa/xla/xla/test# c++filt _ZNK5torch8autograd4Node4nameB5cxx11Evtorch::autograd::Node::name[abi:cxx11]() const

If I install them with r1.13，it's OK。Can you provide a GPU wheel of torch_xla compatible with torch2.0 ，like torch_xla r2.0's wheel ? Or other solutrions.

[JackCaoG]

@qipengh I think this is due to the ABI issue, you can use the torch wheel we built

pip install https://storage.googleapis.com/tpu-pytorch/wheels/cuda/117/torch-nightly+20230228-cp38-cp38-linux_x86_64.whl
pip install https://storage.googleapis.com/tpu-pytorch/wheels/cuda/117/torch_xla-nightly+20230228-cp38-cp38-linux_x86_64.whl

(pretty much just replace torch_xla with torch). For the release wheel we will make sure they are compatible with the upstream wheel.

[qipengh]

Is https://storage.googleapis.com/tpu-pytorch/wheels/cuda/117/torch-nightly+20230228-cp38-cp38-linux_x86_64.whl compiled and built with USE_CUDA=0 ? after install the whl, the error is: can not found libc10_cuda.so. And i need install NVIDIA/apex.

Traceback (most recent call last):
  File "run_squad.py", line 42, in <module>
    from optimization import BertAdam, warmup_linear
  File "/projs/framework/huangqipeng/pytorch/tt/xla/DeepLearningExamples/PyTorch/LanguageModeling/BERT/optimization.py", line 27, in <module>
    import amp_C
ImportError: libc10_cuda.so: cannot open shared object file: No such file or directory

real	0m2.244s
user	0m2.011s
sys	0m0.179s
(pytorch) root@PowerEdge-C4140:/projs/framework/huangqipeng/pytorch/tt/xla/DeepLearningExamples/PyTorch/LanguageModeling/BERT# ldd /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/_XLAC.cpython-38-x86_64-linux-gnu.so
	linux-vdso.so.1 (0x00007ffd96fd3000)
	libc10.so => /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libc10.so (0x00007f1a6f64c000)
	libtorch.so => /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch.so (0x00007f1a6f44a000)
	libtorch_cpu.so => /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (0x00007f1a63b33000)
	libtorch_python.so => /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch_python.so (0x00007f1a62666000)
	libxla_computation_client.so => /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch_xla/lib/libxla_computation_client.so (0x00007f1a38cb9000)
	libstdc++.so.6 => /root/anaconda3/envs/pytorch/lib/libstdc++.so.6 (0x00007f1a70253000)
	libm.so.6 => /lib/x86_64-linux-gnu/libm.so.6 (0x00007f1a3891b000)
	libgcc_s.so.1 => /root/anaconda3/envs/pytorch/lib/libgcc_s.so.1 (0x00007f1a70237000)
	libpthread.so.0 => /lib/x86_64-linux-gnu/libpthread.so.0 (0x00007f1a386fc000)
	libc.so.6 => /lib/x86_64-linux-gnu/libc.so.6 (0x00007f1a3830b000)
	/lib64/ld-linux-x86-64.so.2 (0x00007f1a701a2000)
	libmkl_intel_lp64.so.1 => /root/anaconda3/envs/pytorch/lib/libmkl_intel_lp64.so.1 (0x00007f1a3776c000)
	libmkl_intel_thread.so.1 => /root/anaconda3/envs/pytorch/lib/libmkl_intel_thread.so.1 (0x00007f1a3401d000)
	libmkl_core.so.1 => /root/anaconda3/envs/pytorch/lib/libmkl_core.so.1 (0x00007f1a2fbaf000)
	libomp.so.5 => /usr/lib/x86_64-linux-gnu/libomp.so.5 (0x00007f1a2fac0000)
	libdl.so.2 => /lib/x86_64-linux-gnu/libdl.so.2 (0x00007f1a2f8bc000)
	librt.so.1 => /lib/x86_64-linux-gnu/librt.so.1 (0x00007f1a2f6b4000)
	libshm.so => /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libshm.so (0x00007f1a2f4ac000)

Does the community have a torch whl that matching USE_CUDA=1 and not conflict with undefined symbol reported by torch_xla?

[JackCaoG]

I might be able to kick off a build today to get your that wheel, let me update you later today.

[JackCaoG]

OK let's see if two cloud build cuda-117-nightly-pt38-release in #4728 will pass. if it is you should see some wheel under https://storage.googleapis.com/tpu-pytorch/wheels/cuda/117/use_cuda/

[qipengh]

#4728, pytorch_xla_run_CPU_test of CI is failed. @JackCaoG

[kevint324]

from the Toy Example

import torch
def fn(x, y):
    a = torch.sin(x).cuda()
    b = torch.sin(y).cuda()
    return a + b
new_fn = torch.compile(fn, backend="inductor")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn(input_tensor, input_tensor)

torch.compile(fn, backend="inductor") can be used along with .cuda()

Does below code work if I'm using XLA-GPU?

new_fn = torch.compile(fn, backend="torchxla_trace_once")
input_tensor = torch.randn(10000).to(device="cuda:0")

From the UX point of view, is it a design goal that backend = torchxla_trace_once and backend = inductor are somewhat interchangable under nvidia platform?