Do transfers from PyTorch CUDA to XLA CUDA round trip to CPU?

[Linux-cpp-lisp]

If I have a PyTorch CUDA tensor and am using torch-xla on GPU, will this code roundtrip to CPU?:

t = torch.randn(2, 2, device="cuda")
t2.to(device=xm.xla_device())

Similarly, with:

t = torch.randn(2, 2, device=xm.xla_device())
t = t.square()
t2 = t.cuda()

does t.cuda() (1) trigger graph execution/synchronization and (2) does it involve a CPU round-trip?

(Context: we run models for inference in a context where the input data is provided already on the GPU and put into a torch tensor using torch::from_blob; it would be wasteful and likely eliminate any performance gains from XLA if we had to round trip the input data to CPU.)

Thanks!

[JackCaoG]

@bdhirsh might have a better idea but I am pretty syre cuda tensor will have a round trip to cpu before moving to xla.

for you second question

1. t.cuda() will trigger a graph execution
2. it will round trip to cpu I believe

[bdhirsh]

I'm not 100% sure, but I think it might avoid a round trip to CPU. On the pytorch core side when you call cuda_tensor.to(xla_tensor), we end up eventually calling:

  // create an empty tensor on XLA
  auto r = at::empty(self.sizes(), options.memory_format(memory_format).pinned_memory(pin_out), c10::nullopt);
  // call xla_tensor.copy_(cuda_tensor)  
  r.copy_(self, non_blocking);

That eventually makes its way to at::_copy_from(src=cuda_tensor, dst=xla_tensor), which is implemented in pt/xla (here). Where eventually TensorToXLAData() is called directly on the cuda tensor. I couldn't tell immediately if that involves a roundtrip to CPU or not.

[Linux-cpp-lisp]

I see, thanks for the info— I guess one could check this in a hacky way by just benchmarking:

t = torch.randn(2_000_000, 2, device=xm.xla_device())
t = t.square()
t2 = t.cuda()
torch.cuda.synchronize()

vs

t = torch.randn(2_000_000, 2, device=xm.xla_device())
t = t.square()
t2 = t.cpu()
t2 = t.cuda()
torch.cuda.synchronize()

Probably there's also some way to get nvtop/nvidia-smi/whatever to output bus traffic amounts, but I'm not finding it.

[cicirori]

@Linux-cpp-lisp Hi, did you running this code? I am doing the same test. And it will always trigger a core dump. see #3466

With your demo code got core dump too:

import torch
import torch_xla.core.xla_model as xm

t = torch.randn(2, 2, device="cuda")
t.to(device=xm.xla_device())

[Linux-cpp-lisp]

@cicirori I did not and that's a shame; I've been away from my GPU and wanted to find out more about support for this feature (or lack thereof) before devoting significant resources to compiling torch-xla with CUDA support. (Direct CUDA to XLA-GPU transfers, and TorchScript support, are dealbreakers for me.)

[cicirori]

@cicirori I did not and that's a shame; I've been away from my GPU and wanted to find out more about support for this feature (or lack thereof) before devoting significant resources to compiling torch-xla with CUDA support. (Direct CUDA to XLA-GPU transfers, and TorchScript support, are dealbreakers for me.)

XD, actually I also spent nearly one day to build a torch xla image both torch cuda available and torch xla cuda available (mostly because access websites like GitHub is slow in my dev environment).

@JackCaoG I noticed there are cuda supported images in https://console.cloud.google.com/gcr/images/tpu-pytorch/GLOBAL/xla. But unfortunately PyTorch in the ”cuda“ images isn't compiled with cuda enabled. I wanna know if you guys can enable PyTorch USE_CUDA in your build flows, this would be very helpful for investing problem about GPU backend. Thanks a lot!

(I do have images with cuda support for both torch and torch_xla before building with Github version. But the code of torch xla my images used is slightly different with the repo. So it's very helpful if you guys can bring OFFICIAL ones!)

[JackCaoG]

@cicirori I actually never think of that. We disable PyTorch cuda because we never think of a use case of both pytorch and pytorch/xla use cuda at the same setup. I am not sure if that is allowed(seems like it is actually supported)?

[cicirori]

We disable PyTorch cuda because we never think of a use case of both pytorch and pytorch/xla use cuda at the same setup.

I will encounter many cases where I need to compare the performance of torch cuda and torch xla gpu.

Or maybe I need to do a quick demo on torch cuda (since torch xla compilation is quite long on some models).

By itself, it makes sense to provide an official environment that supports both, for developers using torch xla on gpu devices to develop some features (e.g. fallback to torch cuda other than cpu).

I am not sure if that is allowed(seems like it is actually supported)?

The answer should be yes. I experimented a lot on my own build of the docker image, and both torch cuda and xla gpu will be used (not running at the same time).

[Linux-cpp-lisp]

@JackCaoG: while I second everything @cicirori said, we have an even more direct need for both PyTorch and XLA CUDA, which is running performance-sensitive inference on data that is already on the GPU. Because XLA tensors are opaque, we cannot create an XLA CUDA tensor directly from CUDA memory, but we can (using torch::from_blob) create PyTorch CUDA tensors directly from CUDA memory blobs.

The main reason I'm even looking into torch-xla here is that our models are performance-critical and often overhead dominated (so I hope that the JIT will help), but as a result round-tripping input data to the CPU is an unacceptable performance hit.

For me, of course, all of this also depends on whether torch-xla can be used in C++ with libtorch, which I'm hoping to get some clarity on in another issue: #3451.

Thanks!