diff --git a/modelopt/torch/opt/conversion.py b/modelopt/torch/opt/conversion.py index bb3fe0ce4a5..4cd8e894588 100644 --- a/modelopt/torch/opt/conversion.py +++ b/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 diff --git a/modelopt/torch/opt/mode.py b/modelopt/torch/opt/mode.py index d9ac1bf07bf..50274a4cd91 100644 --- a/modelopt/torch/opt/mode.py +++ b/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 diff --git a/tests/unit/torch/opt/test_chaining.py b/tests/unit/torch/opt/test_chaining.py index 3bc294f3b1e..47b73268de4 100644 --- a/tests/unit/torch/opt/test_chaining.py +++ b/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"]