Add a runner for MosaicML cloud (#44)

* add mosaic launcher string constant

* first attempt at mosaic multinode runner for gptneox

* typo

* actually add mosaic runner

* string cast

* debug

* fix

* actually fix?

* strip extra space

* debugging

* debugging

* correctly set env vars

* drop cd

* add env vars via env instead of export commands

* fix

* try using slurms arg parsing

* debug print

* debug print

* print debug

* cleanup

* try getting world info from the hostfile

* add missing init arg

* more cleanup

* remove more prints

deepspeed/launcher/constants.py

@@ -5,6 +5,7 @@
 
 OPENMPI_LAUNCHER = 'openmpi'
 SLURM_LAUNCHER = 'slurm'
+MOSAICML_LAUNCHER = 'mosaicml'
 
 MVAPICH_LAUNCHER = 'mvapich'
 MVAPICH_TMP_HOSTFILE = '/tmp/deepspeed_mvapich_hostfile'

deepspeed/launcher/multinode_runner.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import os
 import sys
@@ -252,3 +253,40 @@ def get_cmd(self, environment, active_resources):
 
  return mpirun_cmd + export_cmd + python_exec + [self.user_script
  ] + self.user_arguments
+class MosaicMLRunner(MultiNodeRunner):
+ def __init__(self, args, world_info_base64):
+ super().__init__(args, world_info_base64)
+
+ def backend_exists(self):
+ return True
+
+ def parse_user_args(self):
+ user_args = []
+ for arg in self.args.user_args:
+ if arg.startswith('{') and arg.endswith('}'):
+ try:
+ arg_dict = json.loads(arg)
+ if 'config_files' in arg_dict:
+ config_files = {}
+ for k, v in arg_dict.get('config_files', {}).items():
+ config_files[k] = json.loads(v)
+ arg_dict['config_files'] = config_files
+ except json.JSONDecodeError as jde:
+ raise ValueError('Please use plain json for your configs. Check for comments and lowercase trues') from jde
+ arg = json.dumps(arg_dict, separators=(',', ':'))
+ user_args.append(arg)
+ return user_args
+
+ def get_cmd(self, environment, active_resources):
+ deepspeed_launch = [
+ sys.executable,
+ "-u",
+ "-m",
+ "deepspeed.launcher.launch",
+ '--world_info={}'.format(self.world_info_base64),
+ "--node_rank={}".format(os.environ['NODE_RANK']),
+ "--master_addr={}".format(os.environ['MASTER_ADDR']),
+ "--master_port={}".format(os.environ['MASTER_PORT']),
+ ]
+
+ return deepspeed_launch + [self.user_script] + self.user_arguments

deepspeed/launcher/runner.py

@@ -18,8 +18,8 @@
 
 import torch.cuda
 
-from .multinode_runner import PDSHRunner, OpenMPIRunner, MVAPICHRunner, SlurmRunner
-from .constants import PDSH_LAUNCHER, OPENMPI_LAUNCHER, MVAPICH_LAUNCHER, SLURM_LAUNCHER
+from .multinode_runner import PDSHRunner, OpenMPIRunner, MVAPICHRunner, SlurmRunner, MosaicMLRunner
+from .constants import PDSH_LAUNCHER, OPENMPI_LAUNCHER, MVAPICH_LAUNCHER, SLURM_LAUNCHER, MOSAICML_LAUNCHER
 from ..constants import TORCH_DISTRIBUTED_DEFAULT_PORT
 from ..utils import logger
 
@@ -343,6 +343,8 @@ def main(args=None):
  runner = MVAPICHRunner(args, world_info_base64, resource_pool)
  elif args.launcher == SLURM_LAUNCHER:
  runner = SlurmRunner(args, world_info_base64, resource_pool)
+ elif args.launcher == MOSAICML_LAUNCHER:
+ runner = MosaicMLRunner(args, world_info_base64)
  else:
  raise NotImplementedError(f"Unknown launcher {args.launcher}")
 
@@ -367,11 +369,10 @@ def main(args=None):
  for var in fd.readlines():
  key, val = var.split('=')
  runner.add_export(key, val)
-
  cmd = runner.get_cmd(env, active_resources)
 
  logger.info("cmd = {}".format(' '.join(cmd)))
- result = subprocess.Popen(cmd, env=env)
+ result = subprocess.Popen(cmd, env=dict(env, **runner.exports))
  result.wait()
 
  # In case of failure must propagate the error-condition back to the caller (usually shell). The