Update Python device API for SPMD

[JackCaoG]

The purpose of this pr is to clean up the xla_model api call behavior under the spmd context. IMO we should always return xla:0 to the user and use that to represent the global device. The other api like global_device_count and ordinal will also behave as there is only one device. The idea here is that all of the existing api under xla_model is about the xla device. Under SPMD context there is one XLA device, which is a virtual device.

We provided a set of separate api that has runtime in their name for user to query the real runtime device information. for example, user should use

xr. global_runtime_device_count()

instead of

xr. global_device_count()  --> always return 1 when `XLA_USE_SPMD=1`

or

len(xm.get_xla_supported_devices()) --> always return 1 when `XLA_USE_SPMD=1`

This pr will require some model code change, hence I want to push this before next release.

Next step

1. add similar test for apis under runtime.py
2. add a warning message when user using global_device_count(and other similar apis) under SPMD context
3. clean up the SPMD device in C++ land

[will-cromar]

LGTM in principle

I'm on the fence about the naming convention here. So far, I've been thinking about the distinction between a TPU chip (an actual PCI device) and a TPU core (a "device" within that chip) by calling the TPU core a "logical" device and the TPU chip a "physical" device.

We should settle on terms we can use by convention for each of these sets:

1. The set of all runtime devices (e.g. TPU cores)
2. The set of all PCI devices (e.g. TPU chip)
3. The set of all PyTorch devices (either the SPMD virtual device or runtime devices)

To me, "physical" makes the most sense to me as (2), since these are the actual devices you see if you run lspci.

[alanwaketan]

LGTM in principle

I'm on the fence about the naming convention here. So far, I've been thinking about the distinction between a TPU chip (an actual PCI device) and a TPU core (a "device" within that chip) by calling the TPU core a "logical" device and the TPU chip a "physical" device.

We should settle on terms we can use by convention for each of these sets:

1. The set of all runtime devices (e.g. TPU cores)
2. The set of all PCI devices (e.g. TPU chip)
3. The set of all PyTorch devices (either the SPMD virtual device or runtime devices)

To me, "physical" makes the most sense to me as (2), since these are the actual devices you see if you run lspci.

I'm thinking this is very architecture dependent?
For V4, each device is actually corresponding to a chip.
For V3, it's the story you present.

[alanwaketan]

LGTM. Pending Will's comment.

[will-cromar]

Right, (1) and (2) are going to be the same for TPU v4 because the driver represents both cores as one large device. The same is true for GPU.

For TPU v2/v3 and CPU, (1) will be larger than (2) because the runtime will represent the cores as separate devices within the same process.

[JackCaoG]

I am flexibile with the naming convention here, I forgot about the chip vs core for v3... Did we ever expose the chip vs core to user through any api? For example is there any api that returns number of chip which will will return 4 for 'v3-8'?

I think it would be best if user don't need to know about the core vs chips since it is really confusing.

in my mind, under SPMD, we have 1 logic device. In normal data parallel, v3-8 has 8 logical device and v4-8 has 4 logic device.

SPMD

Accelator	logical	physcial
v3-8	1	8
v4-8	1	4

non-SPMD

Accelator	logical	physcial
v3-8	8	8
v4-8	4	4

[will-cromar]

Yeah, there are APIs in tpu.py the differentiate between chips and cores:

xla/torch_xla/_internal/tpu.py

Lines 92 to 119 in 56ae707

	def num_available_chips() -> int:
	"""Returns the number of TPU chips attached through PCI."""
	num_chips = 0
	for vendor_path in glob.glob('/sys/bus/pci/devices/*/vendor'):
	vendor_id = pathlib.Path(vendor_path).read_text().strip()
	if vendor_id != _GOOGLE_PCI_VENDOR_ID:
	continue
	device_path = os.path.join(os.path.dirname(vendor_path), 'device')
	device_id = pathlib.Path(device_path).read_text().strip()
	if device_id in _TPU_PCI_DEVICE_IDS:
	num_chips += 1
	return num_chips
	def num_logical_cores_per_chip() -> int:
	"""Returns number of XLA TPU devices per physical chip on the current host."""
	return 2 if version() <= 3 else 1
	def num_available_devices() -> int:
	"""Returns number of XLA TPU devices on the current host.
	Note: this does not initialize the computation client and is safe to call
	before `xmp.spawn`.
	"""
	return num_available_chips() * num_logical_cores_per_chip()

Concretely, the number of chips is the expected number of local processes when using multiprocessing. I agree that for normal use cases, we should bury this distinction as much as we can. A clear naming convention is important for tools building on us (e.g. Lightning and Accelerate) and for ourselves once we forget this conversation.

I'm fine with "logical" devices being the of torch.devices we expose. What do you think of these?

1. "physical" devices = PCI devices
2. "logical" devices = torch.devices
3. "runtime" devices = PjRtDevices

[JackCaoG]

yea, I am ok with the name runtime devices, I can rename my api to global_runtime_device_count

[JackCaoG]

I am half way editing this pr but now I have a different thought. In https://github.com/pytorch/xla/blob/master/torch_xla/experimental/xla_sharding.py we used term physcial a lot. I also think physical is an easier to understand term, since most people probably don't know what runtime we are using and how runtime device is different than logical device. If we are only going to use the PCI device concept internally(user don't access it), can we just call it pci_device and leave the physical_device to be the number of device that user can create mesh on? @will-cromar

I think runtime_device is a more accurate term, but it would be confusing to user if they don't know how TPU and pytorch/xla runtime works. Physical is more natural IMO.

[JackCaoG]

@will-cromar gentle ping on your opinion toward using physical to represent runtime device.

[JackCaoG]

chatted with Will offline, I think I am just going to use the term defined in #5129 (comment) and update all physcial in this pr to be runtime to be accurate.

[jonb377]

Should these be runtime instead of physical in this file?

[JackCaoG]

good catch

[JackCaoG]

fixed.

[JackCaoG]

@will-cromar @jonb377 Can I get a review for this one?

[jonb377]

Overall LGTM, just a couple of questions to help my understanding

[JackCaoG]

@amithrm @Liyang90 @jonb377 @alanwaketan @baoleai I am going to merge this pr now, the most noticeable change to users are you should not use xr. global_device_count under SPMD since it will return 1(since there is only one xla virtual device). You should use xr. global_runtime_device_count which will tell you the actual number of devices that runtime sees.