[will-cromar]

• Adds experimental implementation of DDP that works on TPU v3 and gives moderately better performance on TPU v4 (1805 ex/sec/replica vs 1710 ex/sec/replica in my experiments).
• Add C++ methods to get PJRT process index and process count (similar to JAX)
• Initialize torch.distributed in pjrt.spawn, since the master address and process ranks are not predictable on TPU.
• Tweak learning rate and atol in DDP tests. The LR was so small that the parameter comparison always passed, even when I had errors in my code.
• Don't broadcast buffers in ResNet50 example to match native XLA distributed implementation.
• Test PJRT with both upstream and experimental DDP.

Tested with PJRT and XRT:

XRT_TPU_CONFIG="localservice;0;localhost:51011" TPU_NUM_DEVICES=4 MASTER_ADDR=localhost MASTER_PORT=12355 python pytorch/xla/test/test_ddp.py
PJRT_DEVICE=TPU python pytorch/xla/test/pjrt/test_ddp.py

[JackCaoG]

one customer asked if they can use torch.distributed to get the ordinal and world size, do you know if that will work if we init the xla backend?

[JackCaoG]

if I understand correctly

1. if user use the upstream ddp, they need to manually init_process_group
2. if user use our custom wrapper, they don't.

If this a fair understanding?

[will-cromar]

Ordinal and world size will be correct in torch.distributed on TPU v4 if the process group is initialized with init_pjrt_process_group, but not on v2/v3, because we'll have two ordinals per rank on v2/v3.

I added the option to initialize the torch.distributed group through xmp.spawn for three reasons, in order of importance:

1. The master IP is not predictable in a TPU slice.
2. If on TPU v2/v3, only one thread can initialize the torch.distributed group. It's easier to do it here than tell the user to manage that themselves.
3. It may be convenient for the PJRT process ID to match each process' rank so we have one fewer index to keep track of.

So, with either DDP implementation, the user can initialize the main process group themselves if they solve the three problems above.

The default is to not initialize the process group since we found that TCP process group initialization is flaky on large pod slices.

[JackCaoG]

hmm, so user don't need to init the process group if they are using our wrapper?

[will-cromar]

See above.

[JackCaoG]

Can we update README and add corresponding tests on xlml? I think we don't have either native ddp nor this new custom ddp e2e test running right now.

[alanwaketan]

LGTM. Thanks, Will. This looks awesome.

[alanwaketan]

I guess we don't need to implement this logic?

[will-cromar]

The broadcast_buffers option is implemented in forward

[alanwaketan]

I guess my confusion is like if the previous optimizer_step approach doesn't demand this, what's the benefits of providing such option here? And then it doesn't seem like you test this option.

[will-cromar]

Filed bugs to update documentation and add tests.