From 63991555ec082c8f80c475f851d008193b10008c Mon Sep 17 00:00:00 2001
From: Hailey Schoelkopf <65563625+haileyschoelkopf@users.noreply.github.com>
Date: Fri, 19 Jan 2024 10:32:17 -0500
Subject: [PATCH] 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 <github-actions@github.com>
---
 configs/neox_arguments.md                 | 19 ++++++++++++++-----
 megatron/neox_arguments/arguments.py      |  6 ++++++
 megatron/neox_arguments/deepspeed_args.py |  7 ++++++-
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/configs/neox_arguments.md b/configs/neox_arguments.md
index a60ef7716..99fe657f7 100644
--- 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
index 8cd5881f7..d293f7b6a 100644
--- 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
index 15b35e411..071c3d3b7 100644
--- 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.
     """