save args

Signed-off-by: zhangsmallshark <zhangsmallshark@gmail.com>

Author: zhangsmallshark

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.")

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: