save args

Signed-off-by: zhangsmallshark <[email protected]>
zhangsmallshark · Feb 12, 2025 · abfedb9 · abfedb9
1 parent b436bd4
commit abfedb9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/training/DeepSpeed-Domino/domino/arguments.py b/training/DeepSpeed-Domino/domino/arguments.py
@@ -218,7 +218,7 @@ def parse_args():
                        help='Do not load optimizer when loading checkpoint.')
     parser.add_argument('--no-load-rng', action='store_true', default=None,
                        help='Do not load rng state when loading checkpoint.')
-    group.add_argument('--exit-on-missing-checkpoint', action='store_true',
+    parser.add_argument('--exit-on-missing-checkpoint', action='store_true',
                        help="If '--load' is set, but checkpoint is not found "
                        "(e.g., path typo), then exit instead of random "
                        "initialization.")

diff --git a/training/DeepSpeed-Domino/megatron/checkpointing.py b/training/DeepSpeed-Domino/megatron/checkpointing.py
@@ -247,7 +247,10 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler):
 
         # Arguments, iteration, and model.
         state_dict = {}
-        # state_dict['args'] = args
+        t_args = args
+        t_args.init_method = None
+        t_args.output_layer_init_method = None
+        state_dict['args'] = t_args
         state_dict['checkpoint_version'] = 3.0
         state_dict['iteration'] = iteration
         if len(model) == 1: