Better distributed init for SLURM

Summary: Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/696

Differential Revision: D16068394

Pulled By: myleott

fbshipit-source-id: 92b44470ab8aeb9f99838cf74e34176104eb2b87

fairseq/distributed_utils.py

@@ -37,6 +37,8 @@ def infer_init_method(args):
 
     # we can determine the init method automatically for Slurm
     elif args.distributed_port > 0:
+        node_list = os.environ.get('SLURM_STEP_NODELIST')
+        if node_list is None:
             node_list = os.environ.get('SLURM_JOB_NODELIST')
         if node_list is not None:
             try: