fix: apply_mode accepts ModeDescriptor instances and an empty mode list with a ModelLike

modelopt/torch/opt/conversion.py

@@ -327,7 +327,8 @@ def __init__(self, modellike: ModelLike) -> None:
     def init_modellike(self) -> nn.Module:
         """Initialize the modellike to be an actual model."""
         model = init_model_from_model_like(self.modellike)
-        ModeloptStateManager.transfer_state_dict(self, model)
+        if ModeloptStateManager.is_converted(self):
+            ModeloptStateManager.transfer_state_dict(self, model)
         return model
 
 

modelopt/torch/opt/mode.py

@@ -356,8 +356,14 @@ def get_registry_by_name(cls, registry_name: str) -> "_ModeRegistryCls":
 
 def get_mode_config(mode_like: ModeLike) -> ModeConfigList:
     """Standardize mode to ModeConfigDict and return."""
-    mode_and_config = [
-        ((m, {}) if isinstance(m, str) else (m[0], m[1] or {})) for m in val2list(mode_like)
-    ]
+
+    def _normalize(m) -> tuple[str, dict]:
+        if isinstance(m, ModeDescriptor):
+            return m.name, {}
+        if isinstance(m, str):
+            return m, {}
+        return m[0], m[1] or {}
+
+    mode_and_config = [_normalize(m) for m in val2list(mode_like)]
 
     return mode_and_config

tests/unit/torch/opt/test_chaining.py

@@ -25,6 +25,8 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
 import modelopt.torch.sparsity as mts
+from modelopt.torch.opt.conversion import ModelLikeModule
+from modelopt.torch.opt.mode import _ModeRegistryCls
 from modelopt.torch.utils.distributed import _serialize
 
 
@@ -196,6 +198,38 @@ def test_model_like_initialization(mode, modellike, expect_exception):
     assert isinstance(model2, torch.nn.Module)
 
 
+def test_apply_mode_with_mode_descriptor_instance():
+    """Test that a ModeDescriptor instance behaves identically to its string name."""
+    model_str = SimpleLinearModel()
+    model_desc = SimpleLinearModel()
+    model_desc.load_state_dict(model_str.state_dict())
+
+    descriptor = _ModeRegistryCls.get_from_any("quantize")
+    model_str = mto.apply_mode(model_str, mode=["quantize"], init_state=True)
+    model_desc = mto.apply_mode(model_desc, mode=[descriptor], init_state=True)
+
+    # modelopt state must be identical
+    manager_str = mto.ModeloptStateManager(model_str)
+    manager_desc = mto.ModeloptStateManager(model_desc)
+    assert torch.equal(_serialize(manager_str.state_dict()), _serialize(manager_desc.state_dict()))
+
+    # model state dict must be identical
+    state_str = model_str.state_dict()
+    state_desc = model_desc.state_dict()
+    assert state_str.keys() == state_desc.keys()
+    for key in state_str:
+        assert torch.equal(state_str[key], state_desc[key])
+
+
+def test_apply_mode_empty_mode_with_model_like():
+    """Test that an empty mode list initializes and returns a plain model from a ModelLike."""
+    model = mto.apply_mode((get_model, (), {}), mode=[])
+    assert isinstance(model, torch.nn.Module)
+    assert not isinstance(model, ModelLikeModule)
+    assert not mto.ModeloptStateManager.is_converted(model)
+    model(get_input())
+
+
 def test_sparse_quantized_module():
     model = get_model()
     modes = ["fastnas", "sparse_magnitude"]