Add handling of account when deepspeed_slurm is set

EleutherAI · Jan 19, 2024 · a300979 · a300979
1 parent 2404145
commit a300979
Showing 1 changed file with 6 additions and 0 deletions.
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)