• 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:

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