Merge --pjrt_distributed flag with --ddp flag. (#5732)

pytorch · Apr 22, 2024 · f228a65 · f228a65
1 parent 49db89d
commit f228a65
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 16 deletions.
diff --git a/test/test_train_mp_imagenet.py b/test/test_train_mp_imagenet.py
@@ -31,11 +31,6 @@
     '--ddp': {
         'action': 'store_true',
     },
-    # Use xla:// init_method instead of env:// for `torch.distributed`.
-    # Required for DDP on TPU v2/v3 when using PJRT.
-    '--pjrt_distributed': {
-        'action': 'store_true',
-    },
     '--profile': {
         'action': 'store_true',
     },
@@ -180,11 +175,8 @@ def _train_update(device, step, loss, tracker, epoch, writer):
 
 
 def train_imagenet():
-  if FLAGS.pjrt_distributed:
+  if FLAGS.ddp:
     dist.init_process_group('xla', init_method='xla://')
-  elif FLAGS.ddp:
-    dist.init_process_group(
-        'xla', world_size=xm.xrt_world_size(), rank=xm.get_ordinal())
 
   print('==> Preparing data..')
   img_dim = get_model_property('img_dim')

diff --git a/test/test_train_mp_mnist.py b/test/test_train_mp_mnist.py
@@ -5,9 +5,6 @@
     '--ddp': {
         'action': 'store_true',
     },
-    '--pjrt_distributed': {
-        'action': 'store_true',
-    },
 }
 
 FLAGS = args_parse.parse_common_options(
@@ -76,11 +73,8 @@ def _train_update(device, step, loss, tracker, epoch, writer):
 
 
 def train_mnist(flags, **kwargs):
-  if flags.pjrt_distributed:
+  if flags.ddp:
     dist.init_process_group('xla', init_method='xla://')
-  elif flags.ddp:
-    dist.init_process_group(
-        'xla', world_size=xm.xrt_world_size(), rank=xm.get_ordinal())
 
   torch.manual_seed(1)