fix args

Signed-off-by: zhangsmallshark <[email protected]>

Author: zhangsmallshark

training/DeepSpeed-Domino/domino/arguments.py

@@ -218,6 +218,11 @@ 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',
+                       help="If '--load' is set, but checkpoint is not found "
+                       "(e.g., path typo), then exit instead of random "
+                       "initialization.")
+
     args = parser.parse_args()
 
     args.rank = int(os.getenv('RANK', '0'))

training/DeepSpeed-Domino/domino/training.py

@@ -116,8 +116,6 @@ def setup_model_and_optimizer(base_model,
     else:
         args.iteration = 0
 
-    args.iteration = 0
-
     return model, optimizer, opt_param_scheduler