From 8283d00e1095f377471b17606b7840ce0c224468 Mon Sep 17 00:00:00 2001 From: arham766 Date: Sun, 5 Jul 2026 18:47:48 -0700 Subject: [PATCH 1/2] fix: apply_mode accepts ModeDescriptor instances and an empty mode list with a ModelLike get_mode_config only handled str/tuple entries, so apply_mode(model, [descriptor_instance]) raised TypeError - despite the docstring and the ModeType = ModeDescriptor | str annotation. Descriptors are now normalized to their name up front (registry descriptors are per-name singletons, so name resolution returns the same instance); str/tuple paths are byte-identical, and get_mode_config's only caller is apply_mode. apply_mode((cls, args, kwargs), []) - the documented 'vanilla case (just initialize+return)' - crashed in init_modellike, which unconditionally transferred never-initialized state. The transfer is now guarded by is_converted, keeping the two vanilla paths consistent: both return a plain initialized model without modelopt state. The state-carrying ModelLike path is unchanged (the manager is constructed before init_modellike runs in normal flows). Part of the findings in #1902 (wave 2). Signed-off-by: arham766 --- modelopt/torch/opt/conversion.py | 3 ++- modelopt/torch/opt/mode.py | 9 ++++++- tests/unit/torch/opt/test_chaining.py | 34 +++++++++++++++++++++++++++ 3 files changed, 44 insertions(+), 2 deletions(-) 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..a86bdfdaa23 100644 --- a/modelopt/torch/opt/mode.py +++ b/modelopt/torch/opt/mode.py @@ -357,7 +357,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) + ( + (m.name, {}) + if isinstance(m, ModeDescriptor) + else (m, {}) + if isinstance(m, str) + else (m[0], m[1] or {}) + ) + 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"] From 54b9a032131ef639b8e9cae6ed7638b9cae92434 Mon Sep 17 00:00:00 2001 From: arham766 Date: Sun, 5 Jul 2026 19:15:11 -0700 Subject: [PATCH 2/2] refactor: explicit branches instead of nested ternary in get_mode_config Signed-off-by: arham766 --- modelopt/torch/opt/mode.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/opt/mode.py b/modelopt/torch/opt/mode.py index a86bdfdaa23..50274a4cd91 100644 --- a/modelopt/torch/opt/mode.py +++ b/modelopt/torch/opt/mode.py @@ -356,15 +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.name, {}) - if isinstance(m, ModeDescriptor) - else (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