Merge --pjrt_distributed flag with --ddp flag.

[will-cromar]

These flags were only separate to support XRT. Now that XRT is gone, these options should always be used together.

[JackCaoG]

is this option tested in the CI?

[will-cromar]

No, but I ran both scripts both ways locally.

[JackCaoG]

maybe we should add it to the GPU CI. just let it run for a couple steps and make sure it works. If we provide a option we might as well test it.

[vanbasten23]

By using -ddp in this command: root@xiowei-gpu:/ansible# PJRT_DEVICE=GPU torchrun --nnodes 1 --nproc-per-node 4 pytorch/xla/test/test_train_mp_imagenet.py --fake_data --batch_size=128 --num_epochs=1 --ddp, the globalRate remains to be 0:

| Training Device=xla:0/1 Epoch=1 Step=320 Loss=0.04107 Rate=0.00 GlobalRate=0.00 Time=21:42:42
| Training Device=xla:0/3 Epoch=1 Step=320 Loss=0.04107 Rate=0.00 GlobalRate=0.00 Time=21:42:42
| Training Device=xla:0/2 Epoch=1 Step=320 Loss=0.04107 Rate=0.00 GlobalRate=0.00 Time=21:42:42
| Training Device=xla:0/0 Epoch=1 Step=320 Loss=0.04107 Rate=0.00 GlobalRate=0.00 Time=21:42:42

Previously, the globalRate is much higher:

| Training Device=xla:0/0 Epoch=1 Step=320 Loss=0.00804 Rate=419.54 GlobalRate=247.63 Time=21:46:35
| Training Device=xla:0/1 Epoch=1 Step=320 Loss=0.00804 Rate=419.51 GlobalRate=247.71 Time=21:46:35
| Training Device=xla:0/2 Epoch=1 Step=320 Loss=0.00804 Rate=419.50 GlobalRate=247.70 Time=21:46:35
| Training Device=xla:0/3 Epoch=1 Step=320 Loss=0.00804 Rate=419.49 GlobalRate=247.69 Time=21:46:35

For torchrun, we need to do dist.init_process_group('xla', init_method='xla://') but don't necessarily do ddp? @will-cromar