Update

Signed-off-by: Jingyu Xin <[email protected]>

Author: jingyu-ml

modelopt/torch/peft/__init__.py

@@ -18,6 +18,7 @@
 from . import mode
 from .config import *
 from .convert import *
+
 # isort: off
 # Import plugins last to avoid circular imports
-# from . import plugins
+from . import plugins

modelopt/torch/peft/conversion.py

@@ -15,26 +15,20 @@
 
 """Quantization conversion/restore utilities."""
 
-import fnmatch
-from collections.abc import Callable
-from contextlib import contextmanager
 from typing import Any
 
 import torch.nn as nn
 
 from modelopt.torch.opt.conversion import ApplyModeError, ModelLikeModule, ModeloptStateManager
-from modelopt.torch.opt.dynamic import _DMRegistryCls
 from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict
 from modelopt.torch.utils import get_unwrapped_name
 
-from .config import (
-    PEFTConfig,
-    _QuantizeExportConfig,
-)
-from .lora.layer import LoRAModuleRegistry
+from .config import PEFTConfig, _QuantizeExportConfig
+from .lora.layer import LoRAModule, LoRAModuleRegistry
 
 __all__ = [
     "replace_lora_module",
+    "update_peft_metadata_in_model",
 ]
 
 
@@ -48,46 +42,88 @@ def convert_to_peft_model(model: ModelLikeModule, config: PEFTConfig) -> Convert
     # set_quantizer_by_cfg(model, config.get("quant_cfg", {}))
 
     metadata = {}
-    # update_quantize_metadata(model, config, metadata)
+    update_peft_metadata(model, config, metadata)
 
     return model, metadata
 
+
 def restore_peft_model(
     model: ModelLikeModule, config: PEFTConfig, metadata: MetadataDict
 ) -> nn.Module:
-    #TODO: implemente the restore logic
-    pass
-
-
-
-def update_peft_metadata(
-    model: nn.Module, config: PEFTConfig, metadata: MetadataDict
-) -> None:
-    """Update the quantizer state in the metadata dict."""
-    pass
-
+    convert_to_peft_model(model, config)
+    return restore_peft_state(model, metadata)
+
+
+def restore_peft_state(model: ModelLikeModule, metadata: MetadataDict):
+    """Restore PEFT state from metadata or extra_state.
+    For backward compatibility, we check metadata first. For distributed
+    checkpoints (NeMo-MCore), the state will be in extra_state of each LoRAModule
+    and will be restored automatically via set_extra_state() during load_state_dict().
+
+    Args:
+        model: Model with LoRA modules to restore
+        metadata: Metadata dictionary that may contain peft_state
+    Returns:
+        The model with restored PEFT state
+    """
+    if "peft_state" not in metadata:
+        # For distributed checkpoints (NeMo-MCore), peft_state is stored
+        # in each LoRAModule's extra_state and will be restored via
+        # set_extra_state() during load_state_dict()
+        return model
+
+    # Legacy path: restore from metadata
+    peft_state_dict = metadata["peft_state"]
+    for name, module in model.named_modules():
+        if isinstance(module, LoRAModule):
+            unwrapped_name = get_unwrapped_name(name)
+            if unwrapped_name in peft_state_dict:
+                try:
+                    module.set_from_peft_state(peft_state_dict[unwrapped_name])
+                except Exception as e:
+                    raise ApplyModeError(f"Failed to restore PEFT state for module {name}: {e}")
+
+    return model
+
+
+def update_peft_metadata(model: nn.Module, config: PEFTConfig, metadata: MetadataDict) -> None:
+    """Update the PEFT/LoRA state in the metadata dict."""
+    metadata["peft_state"] = peft_state(model)
+
+
+def peft_state(model: nn.Module) -> dict[str, Any]:
+    return {
+        get_unwrapped_name(n): m.get_peft_state()
+        for n, m in model.named_modules()
+        if isinstance(m, LoRAModule)
+    }
+
+
+def replace_lora_module(
+    model: nn.Module, version=None, config: PEFTConfig = None, registry=LoRAModuleRegistry
+):
+    """Recursively replace the module with LoRA module."""
+    # Register custom plugins (e.g., for Megatron distributed checkpointing)
+    from .custom import register_custom_model_plugins_on_the_fly
 
-def replace_lora_module(model: nn.Module, version=None, config: PEFTConfig = None, registry=LoRAModuleRegistry):
-    """Recursively replace the module with quantized module."""
-    #TODO: register the extra state for megatron-lm
+    register_custom_model_plugins_on_the_fly(model)
 
     if type(model) in registry:
         model = registry.convert(model)
     _replace_lora_module(model, version=version, registry=registry)
 
+
 def export_peft_model(model: nn.Module, config):
     """Export the quantized model to a quantized model."""
     raise NotImplementedError("Exporting a quantized model is not supported yet.")
 
 
-def restore_export_peft_model(
-    model: nn.Module, config, metadata: MetadataDict
-):
+def restore_export_peft_model(model: nn.Module, config, metadata: MetadataDict):
     """Restores the quantized model from the given state dict."""
     raise NotImplementedError("Restoring a quantized & exported model is not supported yet.")
 
 
-def _replace_lora_module(model: nn.Module, version=None,registry=LoRAModuleRegistry):
+def _replace_lora_module(model: nn.Module, version=None, registry=LoRAModuleRegistry):
     for name, child in model.named_children():
         if type(child) in registry:
             lora_module = registry.convert(child)
@@ -106,3 +142,30 @@ def restore_export_quantized_model(
 ) -> nn.Module:
     """Restores the quantized model from the given state dict."""
     raise NotImplementedError("Restoring a quantized & exported model is not supported yet.")
+
+
+def update_peft_metadata_in_model(model: nn.Module) -> None:
+    """Update the PEFT metadata in the model's ModeloptStateManager.
+    This function should be called after manually modifying LoRA adapters to ensure
+    the metadata stored in the ModeloptStateManager reflects the current state.
+
+    Args:
+        model: Model with LoRA modules whose metadata needs updating
+    Example:
+        >>> # After manually adding/modifying adapters
+        >>> for module in model.modules():
+        ...     if isinstance(module, LoRAModule):
+        ...         module.update_layer_lora("custom_adapter", rank=32)
+        >>> # Update metadata to reflect changes
+        >>> update_peft_metadata_in_model(model)
+    """
+    # Check if model has ModeloptStateManager (has been converted with peft mode)
+    if not ModeloptStateManager.is_converted(model):
+        return
+
+    # Get the state manager
+    manager = ModeloptStateManager(model)
+
+    # Update the metadata with current PEFT state
+    if manager._state and manager._last_metadata is not None:
+        manager._last_metadata["peft_state"] = peft_state(model)

modelopt/torch/peft/convert.py

@@ -16,27 +16,28 @@
 """User-facing quantization API."""
 
 import fnmatch
-import inspect
-import warnings
-from collections.abc import Callable, Iterable
 from typing import Any
 
-import torch
 import torch.nn as nn
 
 # import modelopt.torch.quantization as mtq
 from modelopt.torch.opt import apply_mode
+
+# from modelopt.torch.quantization.conversion import set_quantizer_by_cfg
+from modelopt.torch.opt.conversion import ModeloptStateManager
+
 # from modelopt.torch.opt.searcher import ForwardLoop
 # from modelopt.torch.opt.utils import forward_with_reshard
 from modelopt.torch.peft.config import PEFTConfig
-# from modelopt.torch.quantization.conversion import set_quantizer_by_cfg
+
+from .lora.layer import LoRAModule
 
 # from . import config
 # from .algorithms import AutoQuantizeSearcher
 # from .config import QuantizeAlgoCfgType
 # from .conversion import set_quantizer_attribute
 from .mode import PEFTModeRegistry
-from .lora.layer import LoRAModule
+
 # from .nn import QuantModule, TensorQuantizer
 
 # __all__ = [
@@ -50,17 +51,19 @@
 #     "quantize",
 # ]
 
+
 def update_model(
     model: nn.Module,
     config: dict[str, Any | PEFTConfig],
 ):
-    #TODO: deal with extra state, how to save the model
-    #TODO: sharded dict
-    #TODO: metadate
-    #TODO: how to restore the model
+    # TODO: deal with extra state, how to save the model
+    # TODO: sharded dict
+    # TODO: metadate
+    # TODO: how to restore the model
     apply_mode(model, mode=[("peft", config)], registry=PEFTModeRegistry)
     return add_adapter(model, config)
 
+
 def add_adapter(model, config):
     adapter_cfg = config["adapter_cfg"]
     adapter_name = config["adapter_name"]
@@ -77,4 +80,34 @@ def add_adapter(model, config):
                 else:
                     raise NotImplementedError(f"Unsupported type {type(wildcard_or_filter_func)}")
                 module.update_layer_lora(adapter_name, adapter_setting["rank"])
-    return model
+
+    # Update the metadata in ModeloptStateManager after adding adapters
+    _update_peft_metadata_in_state(model)
+    return model
+
+
+def _update_peft_metadata_in_state(model: nn.Module) -> None:
+    """Update the PEFT metadata in the ModeloptStateManager.
+
+    This function updates the metadata to reflect the current state of LoRA adapters
+    after they have been added or modified.
+    """
+    # Check if model has ModeloptStateManager (has been converted with peft mode)
+    if not ModeloptStateManager.is_converted(model):
+        return
+
+    # Get the state manager
+    manager = ModeloptStateManager(model)
+
+    # Get current PEFT state from all LoRA modules
+    current_peft_state = {}
+    for name, module in model.named_modules():
+        if isinstance(module, LoRAModule):
+            from modelopt.torch.utils import get_unwrapped_name
+
+            unwrapped_name = get_unwrapped_name(name)
+            current_peft_state[unwrapped_name] = module.get_peft_state()
+
+    # Update the metadata in the last mode state (which should be 'peft')
+    if manager._state and manager._last_metadata is not None:
+        manager._last_metadata["peft_state"] = current_peft_state

modelopt/torch/peft/custom.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Custom PEFT/LoRA plugins registry."""
+
+# Registry for custom model plugins
+CUSTOM_MODEL_PLUGINS = set()
+
+
+def register_custom_model_plugins_on_the_fly(model):
+    """Registers custom PEFT/LoRA plugins on the fly.
+
+    This is called before LoRAModule replacement to allow plugins
+    to configure the model (e.g., for distributed checkpointing).
+    """
+    for callback in CUSTOM_MODEL_PLUGINS:
+        callback(model)