reorganize the example dir (#7097)

examples/data_parallel/README.md

@@ -0,0 +1,2 @@
+## Recommendation
+Please consider using `train_resnet_spmd_data_parallel.py` since it uses SPMD internally and are very likely yield better perfomrance.

examples/train_resnet_ddp.py → examples/data_parallel/train_resnet_ddp.py

@@ -1,4 +1,9 @@
+import sys
+import os
+example_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append(example_folder)
 from train_resnet_base import TrainResNetBase
+
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 import torch.optim as optim
@@ -21,4 +26,5 @@ def _mp_fn(index):
 
 
 if __name__ == '__main__':
+  print('consider using train_resnet_spmd_data_parallel.py instead to get better performance')
   xmp.spawn(_mp_fn, args=())

examples/train_resnet_spmd_data_parallel.py → examples/data_parallel/train_resnet_spmd_data_parallel.py

@@ -1,3 +1,7 @@
+import sys
+import os
+example_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append(example_folder)
 from train_resnet_base import TrainResNetBase
 
 import numpy as np

examples/train_resnet_xla_ddp.py → examples/data_parallel/train_resnet_xla_ddp.py

@@ -1,11 +1,17 @@
+import sys
+import os
+example_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append(example_folder)
 from train_resnet_base import TrainResNetBase
+
 import torch_xla.distributed.xla_multiprocessing as xmp
 import torch_xla.core.xla_model as xm
 
 
 class TrainResNetXLADDP(TrainResNetBase):
 
   def run_optimizer(self):
+    # optimizer_step will call `optimizer.step()` and all_reduce the gradident
     xm.optimizer_step(self.optimizer)
 
 
@@ -15,4 +21,5 @@ def _mp_fn(index):
 
 
 if __name__ == '__main__':
+  print('consider using train_resnet_spmd_data_parallel.py instead to get better performance')
   xmp.spawn(_mp_fn, args=())

examples/train_resnet_benchmark.py → examples/debug/train_resnet_benchmark.py

@@ -1,3 +1,7 @@
+import sys
+import os
+example_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append(example_folder)
 from train_resnet_base import TrainResNetBase
 
 import itertools

examples/train_resnet_profile.py → examples/debug/train_resnet_profile.py

@@ -1,6 +1,9 @@
 import os
-
+import os
+example_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append(example_folder)
 from train_resnet_base import TrainResNetBase
+
 import torch_xla.debug.profiler as xp
 
 # check https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md#environment-variables

examples/fsdp/README.md

@@ -0,0 +1,2 @@
+## Recommendation
+Please consider using `train_decoder_only_fsdp_v2.py` since it uses SPMD internally and are very likely yield better perfomrance.

examples/train_decoder_only_fsdp_v2.py → examples/fsdp/train_decoder_only_fsdp_v2.py

@@ -1,3 +1,7 @@
+import sys
+import os
+example_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append(example_folder)
 import decoder_only_model
 from train_decoder_only_base import TrainDecoderOnlyBase
 

examples/train_resnet_fsdp_auto_wrap.py → examples/fsdp/train_resnet_fsdp_auto_wrap.py

@@ -1,3 +1,7 @@
+import sys
+import os
+example_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
+sys.path.append(example_folder)
 from train_resnet_base import TrainResNetBase
 from functools import partial
 
@@ -10,7 +14,6 @@
 from torch_xla.distributed.fsdp.wrap import (size_based_auto_wrap_policy,
                                              transformer_auto_wrap_policy)
 
-
 class TrainResNetXLAFSDP(TrainResNetBase):
 
   def __init__(self):
@@ -49,4 +52,5 @@ def _mp_fn(index):
 
 
 if __name__ == '__main__':
+  print('consider using train_decoder_only_fsdp_v2.py instead to get better performance')
   xmp.spawn(_mp_fn, args=())