Enable passing of --account to srun / SlurmLauncher (#1126)

* add `account` to Deepspeed args * Add handling of `account` when `deepspeed_slurm` is set * Update NeoXArgs docs automatically --------- Co-authored-by: github-actions <[email protected]>
EleutherAI · Jan 19, 2024 · 6399155 · 6399155
1 parent 90f70ff
commit 6399155
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 6 deletions.
diff --git a/configs/neox_arguments.md b/configs/neox_arguments.md
@@ -111,7 +111,7 @@ Logging Arguments
 
 - **git_hash**: str
 
- Default = 31cb364
+ Default = a300979
 
  current git hash of repository
 
@@ -968,7 +968,7 @@ Text Generation arguments
 
 - **prompt_end**: str
 
- Default =
+ Default = 
 
 
  a single prompt's end. Defaults to newline
@@ -1010,7 +1010,7 @@ Text Generation arguments
 
 - **eval_results_prefix**: str
 
- Default =
+ Default = 
 
  prefix to which to save evaluation results - final fp will be {eval_results_prefix}_eval_results_yy-mm-dd-HH-MM.json
 
@@ -1760,7 +1760,7 @@ Args for deepspeed config
 
  Default = None
 
-
+ 
 
 
 
@@ -2051,4 +2051,13 @@ Args for deepspeed runner (deepspeed.launcher.runner).
 
  Default = None
 
- Adds a `--comment` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometime necessary for cluster rules, or so I've heard.
+ Adds a `--comment` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometimes necessary for cluster rules, or so I've heard.
+
+
+
+- **account**: str
+
+ Default = None
+
+ Adds a `--account` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometimes necessary for cluster rules, or so I've heard.
+
diff --git a/megatron/neox_arguments/arguments.py b/megatron/neox_arguments/arguments.py
@@ -505,6 +505,12 @@ def get_deepspeed_main_args(self):
  args_list.extend(
  self.convert_key_value_to_command_line_arg("comment", comment)
  )
+ account = getattr(self, "account")
+ if account:
+ args_list.extend(
+ self.convert_key_value_to_command_line_arg("account", account)
+ )
+
  # master_address = os.environ['SLURM_JOB_NODELIST'].split('\n')[0]
  # args_list.extend(
  # self.convert_key_value_to_command_line_arg('master_addr', master_address)

diff --git a/megatron/neox_arguments/deepspeed_args.py b/megatron/neox_arguments/deepspeed_args.py
@@ -358,5 +358,10 @@ class NeoXArgsDeepspeedRunner(NeoXArgsTemplate):
 
  comment: str = None
  """
- Adds a `--comment` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometime necessary for cluster rules, or so I've heard.
+ Adds a `--comment` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometimes necessary for cluster rules, or so I've heard.
+ """
+
+ account: str = None
+ """
+ Adds a `--account` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometimes necessary for cluster rules, or so I've heard.
  """