address review feedback on #1074

Author: ChenhanYu

modelopt/torch/opt/conversion.py

@@ -311,6 +311,66 @@ def update_last_state_before_save(self, model: nn.Module) -> None:
             last_mode.update_for_save(model, last_config, self._last_metadata)
             self._last_config = last_config
 
+    @staticmethod
+    def validate_modelopt_state(modelopt_state: Any) -> None:
+        """Validate that the loaded object is a valid modelopt state file.
+
+        Args:
+            modelopt_state: The loaded object to validate.
+
+        Raises:
+            TypeError: If the loaded object is not a dictionary or has invalid types for nested fields.
+            ValueError: If the loaded dictionary doesn't have the expected schema for a modelopt state file.
+        """
+        # Validate that the loaded object is a dictionary
+        if not isinstance(modelopt_state, dict):
+            raise TypeError(
+                f"Expected loaded modelopt state to be a dictionary, "
+                f"but got {type(modelopt_state).__name__}. "
+                f"The file may not be a valid modelopt state file."
+            )
+
+        # Validate that the dictionary has the expected keys
+        required_keys = {"modelopt_state_dict", "modelopt_version"}
+        missing_keys = required_keys - set(modelopt_state.keys())
+        if missing_keys:
+            raise ValueError(
+                f"The loaded modelopt state is missing required keys: {missing_keys}. "
+                f"Expected keys: {required_keys}. "
+                f"The file may not be a valid modelopt state file."
+            )
+
+        # Validate that modelopt_state_dict is a list
+        state_dict = modelopt_state["modelopt_state_dict"]
+        if not isinstance(state_dict, list):
+            raise TypeError(
+                f"Expected 'modelopt_state_dict' to be a list, "
+                f"but got {type(state_dict).__name__}. "
+                f"The file may not be a valid modelopt state file."
+            )
+
+        # Validate that each entry in the state_dict is a tuple with 2 elements
+        for i, entry in enumerate(state_dict):
+            if not isinstance(entry, tuple) or len(entry) != 2:
+                raise ValueError(
+                    f"Expected each entry in 'modelopt_state_dict' to be a tuple of length 2, "
+                    f"but entry {i} is {type(entry).__name__} with length {len(entry) if isinstance(entry, (tuple, list)) else 'N/A'}. "
+                    f"The file may not be a valid modelopt state file."
+                )
+            mode_name, mode_state = entry
+            if not isinstance(mode_name, str):
+                raise TypeError(
+                    f"Expected mode name (first element of tuple) to be a string, "
+                    f"but got {type(mode_name).__name__} at entry {i}. "
+                    f"The file may not be a valid modelopt state file."
+                )
+            if not isinstance(mode_state, dict):
+                raise TypeError(
+                    f"Expected mode state (second element of tuple) to be a dictionary, "
+                    f"but got {type(mode_state).__name__} at entry {i}. "
+                    f"The file may not be a valid modelopt state file."
+                )
+
 
 class ApplyModeError(RuntimeError):
     """Error raised when applying a mode to a model fails."""
@@ -532,54 +592,8 @@ def load_modelopt_state(modelopt_state_path: str | os.PathLike, **kwargs) -> dic
     kwargs.setdefault("map_location", "cpu")
     modelopt_state = torch.load(modelopt_state_path, **kwargs)
 
-    # Validate that the loaded object is a dictionary
-    if not isinstance(modelopt_state, dict):
-        raise TypeError(
-            f"Expected loaded modelopt state to be a dictionary, "
-            f"but got {type(modelopt_state).__name__}. "
-            f"The file may not be a valid modelopt state file."
-        )
-    
-    # Validate that the dictionary has the expected keys
-    required_keys = {"modelopt_state_dict", "modelopt_version"}
-    missing_keys = required_keys - set(modelopt_state.keys())
-    if missing_keys:
-        raise ValueError(
-            f"The loaded modelopt state is missing required keys: {missing_keys}. "
-            f"Expected keys: {required_keys}. "
-            f"The file may not be a valid modelopt state file."
-        )
-    
-    # Validate that modelopt_state_dict is a list
-    state_dict = modelopt_state["modelopt_state_dict"]
-    if not isinstance(state_dict, list):
-        raise TypeError(
-            f"Expected 'modelopt_state_dict' to be a list, "
-            f"but got {type(state_dict).__name__}. "
-            f"The file may not be a valid modelopt state file."
-        )
-    
-    # Validate that each entry in the state_dict is a tuple with 2 elements
-    for i, entry in enumerate(state_dict):
-        if not isinstance(entry, tuple) or len(entry) != 2:
-            raise ValueError(
-                f"Expected each entry in 'modelopt_state_dict' to be a tuple of length 2, "
-                f"but entry {i} is {type(entry).__name__} with length {len(entry) if isinstance(entry, (tuple, list)) else 'N/A'}. "
-                f"The file may not be a valid modelopt state file."
-            )
-        mode_name, mode_state = entry
-        if not isinstance(mode_name, str):
-            raise TypeError(
-                f"Expected mode name (first element of tuple) to be a string, "
-                f"but got {type(mode_name).__name__} at entry {i}. "
-                f"The file may not be a valid modelopt state file."
-            )
-        if not isinstance(mode_state, dict):
-            raise TypeError(
-                f"Expected mode state (second element of tuple) to be a dictionary, "
-                f"but got {type(mode_state).__name__} at entry {i}. "
-                f"The file may not be a valid modelopt state file."
-            )
+    # Validate the loaded modelopt state
+    ModeloptStateManager.validate_modelopt_state(modelopt_state)
 
     return modelopt_state