Support loading checkpoint specific tag (#352)

Signed-off-by: Moshe Island <[email protected]> Co-authored-by: Moshe Island <[email protected]>
deepspeedai · Feb 26, 2024 · 31e2584 · 31e2584
1 parent 81d68a3
commit 31e2584
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/megatron/arguments.py b/megatron/arguments.py
@@ -1041,6 +1041,8 @@ def _add_checkpointing_args(parser):
                        help='Do not save current rng state.')
     group.add_argument('--load', type=str, default=None,
                        help='Directory containing a model checkpoint.')
+    group.add_argument('--load-tag', type=str, default=None,
+                       help='Specific checkpoint tag to load. Ignores latest.')
     group.add_argument('--no-load-optim', action='store_true', default=None,
                        help='Do not load optimizer when loading checkpoint.')
     group.add_argument('--no-load-rng', action='store_true', default=None,

diff --git a/megatron/checkpointing.py b/megatron/checkpointing.py
@@ -542,10 +542,11 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
         if args.finetune:
             loaded_dir, state_dict = model[0].load_checkpoint(load_dir,
                 load_module_strict=strict, load_optimizer_states=False,
-                load_lr_scheduler_states=False, load_module_only=True)
+                load_lr_scheduler_states=False, load_module_only=True,
+                tag=args.load_tag)
         else:
             loaded_dir, state_dict = model[0].load_checkpoint(load_dir,
-                load_module_strict=strict)
+                load_module_strict=strict, tag=args.load_tag)
         if loaded_dir is None:
             print_rank_0('WARNING: could not find the metadata file {} '.format(
                 load_dir))