[PJRT] Experimental support for torch.distributed and DDP on TPU v2/v3

[will-cromar]

• Use experimental ThreadLocalWorld to enable multithreading
• Implement a new torch.distributed init_method that uses PJRT runtime parameters and supports multithreading
• torch.distributed "rank" will become the same as our "ordinal", meaning we have one fewer set of indices to track
• Converge the PJRT ImageNet and MNIST tests with the original ones
• Remove experimental pjrt.DistributedDataParallel now that the upstream version works on v3

Performance comparison using ResNet50 with fake data on TPU v3:

Runtime	DDP	Throughput (ex/sec/replica)
XRT	No	418.54
XRT	Yes	395.97
PJRT	No	~640
PJRT	Yes	~565

Example usage:

import torch
import torch.distributed as dist
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.experimental.pjrt_backend
from torch_xla.experimental import pjrt

def _all_gather(index: int):
  dist.init_process_group('xla', init_method='pjrt://')
  t = torch.tensor([index], dtype=torch.int32, device=xm.xla_device())
  output = [torch.zeros_like(t) for _ in range(dist.get_world_size())]
  dist.all_gather(output, t)
  xm.mark_step()
  print(output)

if __name__ == '__main__':
  xmp.spawn(self._all_gather)

Needs rebasing after #4504 merges

Follow-up:

• Update automated tests
• Update PJRT and DDP documentation before release

[JackCaoG]

@will-cromar is this one ready for review?

[will-cromar]

I'll take another pass tomorrow to polish and add some comments, but this should be ready for review.

[alanwaketan]

Curios why we use torch.Generator() instead?

[will-cromar]

The other option would be to wrap these randn calls in a lock and give them a common global seed, but explicitly creating a new generator with the same seed seems clearer to me. I would have done the same for module initialization, but that case doesn't support a custom RNG.

[alanwaketan]

One last comment. Maybe add a comment to the test case to suggest the reasoning behind?

[alanwaketan]

Wonder if we want to parameterized the init_method with env as well?

[will-cromar]

Good idea. Added another test that skips for TPU <= v3, since env:// doesn't work nicely with multithreading.

[alanwaketan]

LGTM.

[JackCaoG]

Thanks!