Fix has_modelopt_state() for native Torch checkpoint format

megatron/post_training/checkpointing.py

@@ -2,15 +2,13 @@
 
 import logging
 from pathlib import Path
-from typing import Dict, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import modelopt.torch.opt as mto
-import torch
 import torch.nn as nn
 from modelopt.torch.opt.plugins import restore_sharded_modelopt_state
 
 from megatron.core import dist_checkpointing
-from megatron.core.dist_checkpointing.strategies.common import COMMON_STATE_FNAME
 from megatron.core.utils import get_torch_version, is_torch_min_version
 from megatron.training import get_args
 from megatron.training.checkpointing import _load_base_checkpoint, load_checkpoint
@@ -21,35 +19,36 @@
 NEMO_WEIGHT_DIR_NAMES = {"model_weights": "model.", "weights": "module."}
 
 
-def has_modelopt_state(checkpoint_path: str, ignore_kd_state: bool = False) -> bool:
-    """Check if modelopt_state folder exists inside the checkpoint path.
+def has_modelopt_state(checkpoint_path: str) -> bool:
+    """Check if modelopt_state folder exists inside the checkpoint.
     Args:
         checkpoint_path: Path to the checkpoint directory
-        ignore_kd_state: If True, ignore the knowledge distillation state
 
     Returns:
-        True if modelopt_state folder exists when ignore_kd_state is False,
-        True if modelopt_state folder exists when ignore_kd_state is True and has only
-        distillation state, False otherwise
+        True if modelopt_state exists, False otherwise
     """
-    load_dir, _ = get_sharded_load_dir(checkpoint_path)
-    if load_dir is None:
-        return False
-    modelopt_state_path = load_dir / "modelopt_state"
-    if not modelopt_state_path.is_dir():
-        return False
-    elif ignore_kd_state:
-        return _has_only_kd_state(modelopt_state_path)
-    else:
-        return True
-
+    args = get_args()
 
-def _has_only_kd_state(modelopt_state_path: Path) -> bool:
-    modelopt_state = torch.load(modelopt_state_path / COMMON_STATE_FNAME, weights_only=False)
-    modes_dict = modelopt_state["modelopt_state_dict"]
-    if len(modes_dict) == 1 and modes_dict[0][0] == "kd_loss":
-        return True
-    return False
+    try:
+        if args.ckpt_format == "torch":
+            # Non-sharded
+            state_dict, _, _ = _load_base_checkpoint(checkpoint_path, rank0=False)
+            if state_dict is None:
+                return False
+            if "modelopt_state" not in state_dict:
+                return False
+            return True
+        else:
+            # Sharded
+            load_dir, _ = get_sharded_load_dir(checkpoint_path)
+            if load_dir is None:
+                return False
+            if not (load_dir / "modelopt_state").is_dir():
+                return False
+            return True
+    except Exception as e:
+        print_rank_0(f"Failed to inspect checkpoint in {checkpoint_path}: {e}")
+        return False
 
 
 def get_sharded_load_dir(load_dir: str) -> Tuple[Union[Path, None], str]:
@@ -89,41 +88,41 @@ def get_sharded_load_dir(load_dir: str) -> Tuple[Union[Path, None], str]:
     return sharded_load_dir, sharded_prefix
 
 
-def load_modelopt_state(load_dir: Optional[str] = None, model: Optional[nn.Module] = None) -> Dict:
+def load_modelopt_state(model: nn.Module, load_dir: Optional[str] = None) -> None:
     """Loading modelopt_state without loading the model.
 
-    If --use-dist-ckpt, we try to load from the sharded modelopt_state. This will not load the model
-    state_dict. Otherwise, if the checkpoint is not sharded, we load the base checkpoint (that
-    contains the model state as well) and extract the modelopt_state.
+    If distributed checkpointing in use, we try to load from the sharded modelopt_state. This will not
+    load the model state_dict. Otherwise, if the checkpoint is not sharded, we load the base checkpoint
+    (which contains the model state as well) and extract the modelopt_state.
 
     Args:
+        model: the model to load the modelopt_state into
         load_dir: optionally provide a different loading path
-        model: required when loading a sharded checkpoint
     """
     args = get_args()
+    load_dir = load_dir or args.load
 
-    if load_dir is None:
-        load_dir = args.load
-
-    if args.use_dist_ckpt:
-        assert model is not None, "`model` argument required when `args.use_dist_ckpt is True`"
-        sharded_load_dir, _ = get_sharded_load_dir(load_dir)
-        if sharded_load_dir is None:
-            print_rank_0("No sharded checkpoint found. Skipping loading modelopt_state.")
-            return {}
-        restore_sharded_modelopt_state([model], sharded_load_dir)
-    else:
+    if args.ckpt_format == "torch":
+        # Non-sharded
         print_rank_0(f"Loading ModelOpt state from base checkpoint ({load_dir})")
         try:
             state_dict, _, _ = _load_base_checkpoint(args.load, rank0=False)
         except Exception:
             print_rank_0("Failed to load base checkpoint via megatron _load_base_checkpoint!")
+            return
         if state_dict is None:
             print_rank_0("No checkpoint state_dict found. Skipping loading ModelOpt state.")
         else:
             modelopt_state = state_dict.get("modelopt_state", None)
         if modelopt_state is not None:
             mto.restore_from_modelopt_state(model, modelopt_state)
+    else:
+        # Sharded
+        sharded_load_dir, _ = get_sharded_load_dir(load_dir)
+        if sharded_load_dir is None:
+            print_rank_0("No sharded checkpoint found. Skipping loading modelopt_state.")
+            return
+        restore_sharded_modelopt_state([model], sharded_load_dir)
 
 
 def load_modelopt_checkpoint(