Xla: why xla_model.xla_device() so slow at GPU backend?

[qianzhang613]

❓ Questions and Help

it is take 3min to execute xla_model.xla_device(), so what's extra pytorch xla do compare to tensorflow ? i think xla_model.xla_device() will init the tensorflow session whereas the tensorflow init is just some seconds.

thanks

[dlibenzi]

Seems like there is something wrong with your own build.

import time
import torch
import torch_xla
import torch_xla.core.xla_model as xm

s = time.time()
d = xm.xla_device()
t = time.time() - s
print(t)

(pytorch) dlibenzi@davide-dev:~/pytorch/xla$ GPU_NUM_DEVICES=4 python ~/tmp/xdev_speed.py 
6.864998817443848
(pytorch) dlibenzi@davide-dev:~/pytorch/xla$ GPU_NUM_DEVICES=4 python ~/tmp/xdev_speed.py 
7.732343673706055

And this with 4 devices.

[cjolivier01]

I actually experienced this as well the other day. I didn't have time to investigate, but if it continues, I can try to reproduce again (I just turned off cuda). The delay was in the DeviceFactory::AddDevices() call when creating the local grpc session. If I remember correctly, it was in DeviceFactory::AddDevices() call on line 190 of grpc_server_lib.cc -- I didn't have debug mode enabled for the stuff below this so I didn't look into it further.

  if (opts.local_device_mgr == nullptr) {
    std::vector<std::unique_ptr<Device>> devices;
    TF_RETURN_IF_ERROR(
        DeviceFactory::AddDevices(sess_opts, name_prefix, &devices));
    worker_env_.device_mgr = new StaticDeviceMgr(std::move(devices));
    owned_device_manager_.reset(worker_env_.device_mgr);
  } else {
    worker_env_.device_mgr = opts.local_device_mgr;
    owned_device_manager_.reset(nullptr);
  }

[cjolivier01]

My first guess is it's something messed up with my cuda driver, maybe a driver<->cuda<->os version mismatch, i was going to look into it when I had a down-day (if that day every comes) :)

[dlibenzi]

If happens again, while stuck, in another terminal, run:

$ ps ax | grep ptxas

[byronyi]

It usually means that the local TF (or libxla_computation_client.so) you use are built with mismatched CUDA compute capabilities, e.g. built with 3.5,6.0 but running on V100 (7.0) hardware. It won't be a hard failure, but the CUDA driver will compile just in time for all the PTX code, which could take up to an hour in some cases.

[cjolivier01]

That worked for me, when I explicitly set it for 6.2, no delay.

[qianzhang613]

thanks for all, i have a try that rebuild with export TF_CUDA_COMPUTE_CAPABILITIES="6.0,6.1,7.5"

[qianzhang613]

thanks for all, i have a try that rebuild with export TF_CUDA_COMPUTE_CAPABILITIES="6.0,6.1,7.5"

that works ! many thanks.