xm.all_to_all working for TPU Pod ?

[shz0116]

Hi,

I am trying to learn how xm.all_to_all works on a TPU Pod (across multiple nodes).
I see examples using xmp.spawn to fork 1 or 8 processes on a node (xla/test/test_mp_all_to_all.py).
I assume torch_xla.distributed.xla_multiprocessing only work on a TPU node, not for POD, right ?
Is there an example of using xm.all_to_all on a POD, such as TPUv3-128 ?
Do I need to use cluster ? Thanks.

[taylanbil]

xmp.spawn and xmp are for one node, not pods, yes. You would need to use torch_xla.distributed.xla_dist as in this tutorial to run all_to_all (or anything) on a pod.

[shz0116]

@taylanbil Thanks. Could you check #2413 ? Seems still some problem with all_to_all. The data order in results seems not right.

[JackCaoG]

I am able to repo this, trying to see if this is related to the device assigniment we set is incorrect.

[jysohn23]

It comes down to the torch context devices (ex. TPU:0) passed down to the creation of the XRTComputation being lexicographically sorted instead of sorting by the numerical ordinals. Since there's no explicit device list passed into XLATensor::SynTensorsGraph we use the XrtComputationClient::GetReplicationDevices(). This is setup at xmp.spawn() time and it looks like the order of the replication devices set is lexicographically sorted:

2020-11-13 22:02:09 10.128.15.223 [0] torch_xla._XLAC._xla_get_all_devices(): ['CPU:0', 'CPU:1', 'CPU:2', 'CPU:3', 'TPU:0', 'TPU:1', 'TPU:10', 'TPU:11', 'TPU:12', 'TPU:13', 'TPU:14', 'TPU:15', 'TPU:16', 'TPU:17', 'TPU:18', 'TPU:19', 'TPU:2', 'TPU:20', 'TPU:21', 'TPU:22', 'TPU:23', 'TPU:24', 'TPU:25', 'TPU:26', 'TPU:27', 'TPU:28', 'TPU:29', 'TPU:3', 'TPU:30', 'TPU:31', 'TPU:4', 'TPU:5', 'TPU:6', 'TPU:7', 'TPU:8', 'TPU:9']

Looking at the options_.global_device_map:

2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958640: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[CPU:0] = /job:c_tpu_worker/replica:0/task:0/device:XLA_CPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958652: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[CPU:1] = /job:c_tpu_worker/replica:0/task:1/device:XLA_CPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958658: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[CPU:2] = /job:c_tpu_worker/replica:0/task:2/device:XLA_CPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958663: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[CPU:3] = /job:c_tpu_worker/replica:0/task:3/device:XLA_CPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958668: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:0] = /job:c_tpu_worker/replica:0/task:0/device:TPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958674: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:1] = /job:c_tpu_worker/replica:0/task:0/device:TPU:1
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958680: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:10] = /job:c_tpu_worker/replica:0/task:1/device:TPU:2
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958685: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:11] = /job:c_tpu_worker/replica:0/task:1/device:TPU:3
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958690: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:12] = /job:c_tpu_worker/replica:0/task:1/device:TPU:4
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958695: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:13] = /job:c_tpu_worker/replica:0/task:1/device:TPU:5
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958700: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:14] = /job:c_tpu_worker/replica:0/task:1/device:TPU:6
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958706: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:15] = /job:c_tpu_worker/replica:0/task:1/device:TPU:7
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958711: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:16] = /job:c_tpu_worker/replica:0/task:2/device:TPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958716: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:17] = /job:c_tpu_worker/replica:0/task:2/device:TPU:1
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958721: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:18] = /job:c_tpu_worker/replica:0/task:2/device:TPU:2
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958726: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:19] = /job:c_tpu_worker/replica:0/task:2/device:TPU:3
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958737: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:2] = /job:c_tpu_worker/replica:0/task:0/device:TPU:2
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958743: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:20] = /job:c_tpu_worker/replica:0/task:2/device:TPU:4
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958748: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:21] = /job:c_tpu_worker/replica:0/task:2/device:TPU:5
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958753: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:22] = /job:c_tpu_worker/replica:0/task:2/device:TPU:6
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958758: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:23] = /job:c_tpu_worker/replica:0/task:2/device:TPU:7
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958763: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:24] = /job:c_tpu_worker/replica:0/task:3/device:TPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958768: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:25] = /job:c_tpu_worker/replica:0/task:3/device:TPU:1
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958772: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:26] = /job:c_tpu_worker/replica:0/task:3/device:TPU:2
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958777: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:27] = /job:c_tpu_worker/replica:0/task:3/device:TPU:3
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958782: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:28] = /job:c_tpu_worker/replica:0/task:3/device:TPU:4
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958787: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:29] = /job:c_tpu_worker/replica:0/task:3/device:TPU:5
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958794: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:3] = /job:c_tpu_worker/replica:0/task:0/device:TPU:3
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958800: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:30] = /job:c_tpu_worker/replica:0/task:3/device:TPU:6
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958805: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:31] = /job:c_tpu_worker/replica:0/task:3/device:TPU:7
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958810: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:4] = /job:c_tpu_worker/replica:0/task:0/device:TPU:4
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958815: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:5] = /job:c_tpu_worker/replica:0/task:0/device:TPU:5
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958821: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:6] = /job:c_tpu_worker/replica:0/task:0/device:TPU:6
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958827: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:7] = /job:c_tpu_worker/replica:0/task:0/device:TPU:7
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958832: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:8] = /job:c_tpu_worker/replica:0/task:1/device:TPU:0
2020-11-13 22:00:51 10.128.15.223 [0] 2020-11-13 22:00:51.958837: I    4479 tensorflow/compiler/xla/xla_client/computation_client.cc:119] options->global_device_map[TPU:9] = /job:c_tpu_worker/replica:0/task:1/device:TPU:1

The problem is basically with how the global_device_map is populated and ordered by default since the key is a string. Will send out a PR with fix.

[shz0116]

Great.