Add fallback for SLURM config

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

Differential Revision: D16417582

Pulled By: myleott

fbshipit-source-id: 6b4289febcf9290452bb91f1f2181a02c09c82a7

fairseq/distributed_utils.py

@@ -48,7 +48,14 @@ def infer_init_method(args):
                     port=args.distributed_port,
                 )
                 nnodes = int(os.environ.get('SLURM_NNODES'))
-                ntasks_per_node = int(os.environ.get('SLURM_NTASKS_PER_NODE'))
+                ntasks_per_node = os.environ.get('SLURM_NTASKS_PER_NODE')
+                if ntasks_per_node is not None:
+                    ntasks_per_node = int(ntasks_per_node)
+                else:
+                    ntasks = int(os.environ.get('SLURM_NTASKS'))
+                    nnodes = int(os.environ.get('SLURM_NNODES'))
+                    assert ntasks % nnodes == 0
+                    ntasks_per_node = int(ntasks / nnodes)
                 if ntasks_per_node == 1:
                     assert args.distributed_world_size % nnodes == 0
                     gpus_per_node = args.distributed_world_size // nnodes