add local rank explicitly for mpirun (deepspeedai#72)

Co-authored-by: Shuai Zheng <[email protected]>

Author: szhengac

BingBertSquad/nvidia_run_squad_deepspeed.py

@@ -742,6 +742,7 @@ def main():
     parser = get_argument_parser()
 
     deepspeed.init_distributed(dist_backend='nccl')
+    args.local_rank = int(os.environ['LOCAL_RANK'])
 
     # Include DeepSpeed configuration arguments
     parser = deepspeed.add_config_arguments(parser)

bing_bert/deepspeed_train.py

@@ -391,6 +391,7 @@ def prepare_optimizer_parameters(args, model):
 def prepare_model_optimizer(args):
     # Initialize torch distributed
     deepspeed.init_distributed(dist_backend='nccl')
+    args.local_rank = int(os.environ['LOCAL_RANK'])
 
     # Loading Model
     model = BertMultiTask(args)

pipeline_parallelism/train.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import os
 import argparse
 
 import torch
@@ -148,8 +149,9 @@ def train_pipe(args, part='parameters'):
 if __name__ == '__main__':
     args = get_args()
 
-    torch.cuda.set_device(args.local_rank)
     deepspeed.init_distributed(dist_backend=args.backend)
+    args.local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(args.local_rank)
 
     if args.pipeline_parallel_size == 0:
         train_base(args)