Multiprocessing with GPUs setup for workers has errors

[aingo03304]

🐛 Bug

In xla_multiprocessing.py, _parse_workers_config returns OrderedDict.

xla/torch_xla/distributed/xla_multiprocessing.py

Lines 54 to 63 in 9917e41

	def _parse_workers_config(config):
	# XRT_WORKERS='worker:0;ismz9:25822'
	workers = collections.OrderedDict()
	for worker in config.split('|'):
	m = re.match(r'(\w+):(\d+);((grpc://)?[a-zA-Z0-9_\-\.]+:\d+)', worker)
	if not m:
	raise ValueError('Bad worker syntax: {}'.format(worker))
	workers['{}:{}'.format(m.group(1), m.group(2))] = WorkerConfigEntry(
	worker_name=m.group(1), ordinal=int(m.group(2)), host_port=m.group(3))
	return workers

However

xla/torch_xla/distributed/xla_multiprocessing.py

Lines 156 to 157 in 9917e41

	for h, worker in enumerate(wcfg):
	m = re.match(r'(.*):(\d+)$', worker.host_port)

this code tries to access worker.host_port and this raises the error which can be fixed by replacing worker.host_port to wcfg[worker].host_port.

And,

xla/torch_xla/distributed/xla_multiprocessing.py

Lines 163 to 165 in 9917e41

	workers.append('{}:{};grpc://{}:{}'.format(worker.worker_name, gindex,
	m.group(1),
	int(m.group(2)) + i))

this code is appending '{}:{};grpc://{}:{}' but '{}:{};{}:{}' is correct based on the following configuration in CI, because m.group(1) includes 'grpc://'.

xla/.circleci/test.sh

Line 11 in 3eaee46

export XRT_WORKERS="localservice:0;grpc://localhost:40934"

To Reproduce

Steps to reproduce the behavior:

1. With GPUs, run test_train_mp_mnist.py

Expected behavior

Environment

• Reproducible on XLA backend [CPU/TPU]: GPUs
• torch_xla version: master

Additional context

[JackCaoG]

@aingo03304 Thanks for reporting, will take a look. Do you mind sharing the config you used to setup the multi-GPU?

[aingo03304]

@JackCaoG
Thank you for the response. Here are my env variables I set.

export XRT_WORKERS="localservice:0;grpc://localhost:40934"
export XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0|GPU:0;/job:localservice/replica:0/task:0/device:XLA_GPU:0|GPU:1;/job:localservice/replica:0/task:0/device:XLA_GPU:1|GPU:2;/job:localservice/replica:0/task:0/device:XLA_GPU:2|GPU:3;/job:localservice/replica:0/task:0/device:XLA_GPU:3"

By the way, do you have any plan to officially support pytorch-xla for gpu?

[JackCaoG]

HI @aingo03304 , sorry for the late reply. I think I need more time to look into this. For me if only exporting

export XRT_WORKERS="localservice:0;grpc://localhost:40934"
export XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0|GPU:0;/job:localservice/replica:0/task:0/device:XLA_GPU:0|GPU:1;/job:localservice/replica:0/task:0/device:XLA_GPU:1|GPU:2;/job:localservice/replica:0/task:0/device:XLA_GPU:2|GPU:3;/job:localservice/replica:0/task:0/device:XLA_GPU:3"

will hit

xla/torch_xla/distributed/xla_multiprocessing.py

Line 382 in 9917e41

return _run_direct(fn, args, nprocs, join, daemon, start_method)

and everything seems fine.

If I use

export GPU_NUM_DEVICES=4

I will hit a different issue with my setup. I need to look into it a bit more. I don't think you need to setup XRT_DEVICE_MAP and XRT_WORKERS if you export the GPU_NUM_DEVICES. Those should be setup automatically.

[aingo03304]

@JackCaoG Yes, I can reproduce your comment. Thank you for the reply.
The actual situation I mentioned was

export GPU_NUM_DEVICES=4
export XRT_WORKERS="localservice:0;grpc://localhost:40934"
export XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0|GPU:0;/job:localservice/replica:0/task:0/device:XLA_GPU:0|GPU:1;/job:localservice/replica:0/task:0/device:XLA_GPU:1|GPU:2;/job:localservice/replica:0/task:0/device:XLA_GPU:2|GPU:3;/job:localservice/replica:0/task:0/device:XLA_GPU:3"

Then you can reproduce the python error like this.

Traceback (most recent call last):
  File "test_train_mp_mnist.py", line 196, in <module>
    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=FLAGS.num_cores)
  File "/home/aingo03304/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla-1.6-py3.6-linux-x86_64.egg/torch_xla/distributed/xla_multiprocessing.py", line 384, in spawn
    pf_cfg = _pre_fork_setup(nprocs)
  File "/home/aingo03304/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla-1.6-py3.6-linux-x86_64.egg/torch_xla/distributed/xla_multiprocessing.py", line 210, in _pre_fork_setup
    _setup_workers(num_devices)
  File "/home/aingo03304/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla-1.6-py3.6-linux-x86_64.egg/torch_xla/distributed/xla_multiprocessing.py", line 157, in _setup_workers
    m = re.match(r'(.*):(\d+)$', worker.host_port)
AttributeError: 'str' object has no attribute 'host_port'

Clearly, this is an error on the python side.

However, now I encounter an error like this after I fixed the python side code. It seems something that is in XLA and/or NCCL.

2021-02-05 14:47:38.663029: W tensorflow/core/framework/op_kernel.cc:1763] OP_REQUIRES failed at xrt_execute_op.cc:430 : Internal: tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc:105: NCCL operation ncclGroupEnd() failed: unhandled cuda error
2021-02-05 14:47:38.663211: W tensorflow/core/framework/op_kernel.cc:1763] OP_REQUIRES failed at xrt_execute_op.cc:430 : Internal: tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc:105: NCCL operation ncclGroupEnd() failed: unhandled cuda error
2021-02-05 14:47:38.678277: E tensorflow/compiler/xla/xla_client/xla_util.cc:76] >>> Dumping Computation 0
(...omitted...)
Recent warning and error logs:
  OP_REQUIRES failed at xrt_execute_op.cc:430 : Internal: tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc:105: NCCL operation ncclGroupEnd() failed: unhandled cuda error

Any thought on this?

[JackCaoG]

Hi @aingo03304 Sorry for the late reply, the last error seem to coming from the NCCL group config, in https://github.com/tensorflow/tensorflow/blob/b681785237799530623e94455751a02c43eca1d7/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc#L164

I will try to repo this when I am a bit more free, but this error might not be directly related to pt/xla.

@tyoc213 You were also looking into using XLA:GPU locally and made it work. Did you encounter this error before?

[aingo03304]

@JackCaoG I want to fix this and run my code using multiprocessing asap.
I am not that familiar with NCCL API though, can you elaborate on NCCL group config?
In GPU xla backend code you mentioned, I cannot find any reason to fail ncclGroupEnd.

Don't know what ncclUnhandledCudaError exactly means, but the documentation of NCCL describes it as A call to a CUDA function failed.. Is there any way to debug cudaError in tensorflow or pt/xla configuration?

[zcain117]

Is this a setup with multiple machines, each machine having a few GPUs? Or is this 1 machine with 4 GPUs?

If it's single-node, multi GPU you can just use e.g.:

              unset XRT_TPU_CONFIG
              export GPU_NUM_DEVICES=4
              python3 pytorch/xla/test/test_train_mp_mnist.py

this will run multiprocessing with 4 processes

you would need CUDA installed, for example you could use our public docker image: gcr.io/xl-ml-test/pytorch-xla:nightly_3.6_cuda

[aingo03304]

@zcain117 Thank you for the response. My setup is on a single machine.
I'll try with the docker image and let you know the result. Probably my local CUDA setting is somehow improper for pt/xla.

[tyoc213]

Hi @aingo03304 Sorry for the late reply, the last error seem to coming from the NCCL group config, in https://github.com/tensorflow/tensorflow/blob/b681785237799530623e94455751a02c43eca1d7/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc#L164

I will try to repo this when I am a bit more free, but this error might not be directly related to pt/xla.

@tyoc213 You were also looking into using XLA:GPU locally and made it work. Did you encounter this error before?

Sorry for late answer have been a little busy, yeah I can work locally with pytorch/XLA on GPU, but I only have 1 GPU the nearest thing I found was this #2670 @aingo03304 dont know if that patch could help.

[aingo03304]

@zcain117

you would need CUDA installed, for example you could use our public docker image: gcr.io/xl-ml-test/pytorch-xla:nightly_3.6_cuda

It seems that the image is not public...

[aingo03304]

@tyoc213 Actually I fixed the code and can run test_train_mnist.py. And now I encounter NCCL related error above.
(Thank you for the blog post. It really was really helpful!)

@zcain117

If it's single-node, multi GPU you can just use e.g.:

              unset XRT_TPU_CONFIG
              export GPU_NUM_DEVICES=4
              python3 pytorch/xla/test/test_train_mp_mnist.py

this will run multiprocessing with 4 processes

I encounter the same error:

Recent warning and error logs:
  OP_REQUIRES failed at xrt_execute_op.cc:430 : Internal: tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc:105: NCCL operation ncclGroupEnd() failed: unhandled cuda error

Is there any information I should give you so you can figure out what the problem is?

[zcain117]

Sorry about that, I gave the wrong GCR image. The public one should be:

gcr.io/tpu-pytorch/xla:nightly_3.6_cuda ( or you can use nightly_3.7_cuda )

[aingo03304]

@zcain117 I reproduced the same error on docker image you gave.

OMP: Warning #190: Forking a process while a parallel region is active is potentially unsafe.
warning: Linking two modules of different target triples: '/usr/local/cuda-10.1/nvvm/libdevice/libdevice.10.bc' is 'nvptx64-nvidia-gpulibs' whereas '' is 'nvptx64-nvidia-cuda'

warning: Linking two modules of different target triples: '/usr/local/cuda-10.1/nvvm/libdevice/libdevice.10.bc' is 'nvptx64-nvidia-gpulibs' whereas '' is 'nvptx64-nvidia-cuda'

2021-02-16 06:25:11.307955: W     734 tensorflow/core/framework/op_kernel.cc:1763] OP_REQUIRES failed at xrt_execute_op.cc:430 : Internal: tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc:105: NCCL operation ncclGroupEnd() failed: unhandled cuda error
2021-02-16 06:25:11.308034: W     733 tensorflow/core/framework/op_kernel.cc:1763] OP_REQUIRES failed at xrt_execute_op.cc:430 : Internal: tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc:105: NCCL operation ncclGroupEnd() failed: unhandled cuda error

This result is reproduced on another machine with two TitanXp.