address comments

Author: Your Name

modelopt/torch/opt/dynamic.py

@@ -1273,7 +1273,8 @@ def config(self, configurable: bool | None = None) -> dict[str, Any]:
             A dict of ``(parameter_name, choice)`` that specifies an active subnet.
         """
         return {
-            get_unwrapped_name(name): hp.active for name, hp in self.named_hparams(configurable)
+            get_unwrapped_name(name, self): hp.active
+            for name, hp in self.named_hparams(configurable)
         }
 
     def select(self, config: dict[str, Any], strict: bool = True) -> None:

modelopt/torch/opt/plugins/peft.py

@@ -90,7 +90,7 @@ def _new_load_adapter(self, model_id, adapter_name, *args, **kwargs):
         )
         for name, module in self.named_modules():
             if isinstance(module, TensorQuantizer):
-                module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name)])
+                module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name, self)])
 
     return outputs
 

modelopt/torch/quantization/conversion.py

@@ -123,12 +123,12 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata:
 
     for name, module in model.named_modules():
         if isinstance(module, TensorQuantizer):
-            name = get_unwrapped_name(name)
+            name = get_unwrapped_name(name, model)
             module.set_from_modelopt_state(quantizer_state_dict[name])
 
     for name, module in model.named_modules():
         if isinstance(module, QuantModule):
-            name = get_unwrapped_name(name)
+            name = get_unwrapped_name(name, model)
             module.modelopt_post_restore(name)
 
     return model
@@ -166,7 +166,7 @@ def update_quantize_metadata(
 def quantizer_state(model: nn.Module) -> dict[str, Any]:
     """Returns the quantizer state dict describing the quantizer states in the model."""
     return {
-        get_unwrapped_name(n): m.get_modelopt_state()
+        get_unwrapped_name(n, model): m.get_modelopt_state()
         for n, m in model.named_modules()
         if isinstance(m, (TensorQuantizer, SequentialQuantizer))
     }

modelopt/torch/quantization/utils.py

@@ -452,7 +452,7 @@ def get_quantizer_state_dict(model: nn.Module):
     quantizer_state_dict = {}
     for name, module in model.named_modules():
         if isinstance(module, TensorQuantizer):
-            quantizer_state_dict[get_unwrapped_name(name)] = module.state_dict()
+            quantizer_state_dict[get_unwrapped_name(name, model)] = module.state_dict()
     return quantizer_state_dict
 
 
@@ -461,5 +461,8 @@ def set_quantizer_state_dict(model: nn.Module, quantizer_state_dict: dict):
     from .nn import TensorQuantizer
 
     for name, module in model.named_modules():
-        if isinstance(module, TensorQuantizer) and get_unwrapped_name(name) in quantizer_state_dict:
-            module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name)])
+        if (
+            isinstance(module, TensorQuantizer)
+            and get_unwrapped_name(name, model) in quantizer_state_dict
+        ):
+            module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name, model)])