[1/N] ModelOPT PEFT mode support for the megatron-lm

.github/CODEOWNERS

@@ -22,6 +22,7 @@ modelopt/torch/distill @NVIDIA/modelopt-torch-distill-codeowners
 modelopt/torch/export @NVIDIA/modelopt-torch-export-codeowners
 modelopt/torch/nas @NVIDIA/modelopt-torch-nas-prune-codeowners
 modelopt/torch/opt @NVIDIA/modelopt-torch-opt-codeowners
+modelopt/torch/peft @NVIDIA/modelopt-torch-peft-codeowners
 modelopt/torch/prune @NVIDIA/modelopt-torch-nas-prune-codeowners
 modelopt/torch/quantization @NVIDIA/modelopt-torch-quantization-codeowners
 modelopt/torch/sparsity @NVIDIA/modelopt-torch-sparsity-codeowners

modelopt/torch/peft/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+
+"""PEFT API subpackage for torch."""
+
+from . import mode
+from .config import *
+from .conversion import *
+from .convert import *

modelopt/torch/peft/config.py

@@ -0,0 +1,159 @@
+# 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.
+
+"""Configuration classes for PEFT methods."""
+
+from collections.abc import Callable
+
+from pydantic import field_validator
+
+from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField
+
+__all__ = ["ExportPEFTConfig", "PEFTAttributeConfig", "PEFTConfig"]
+
+
+class PEFTAttributeConfig(ModeloptBaseConfig):
+    """Configuration for PEFT adapter attributes."""
+
+    enable: bool = ModeloptField(
+        default=True,
+        title="Enable adapter",
+        description="If True, enables the adapter. If False, by-passes the adapter.",
+    )
+
+    rank: int = ModeloptField(
+        default=64,
+        title="LoRA rank",
+        description=(
+            "The rank (dimension) of the LoRA matrices. "
+            "Higher rank allows more expressiveness but uses more memory."
+        ),
+    )
+
+    scale: float = ModeloptField(
+        default=1.0,
+        title="LoRA scaling factor",
+        description="Scaling factor for the LoRA output. Controls the magnitude of the adaptation.",
+    )
+
+    lora_a_init: str = ModeloptField(
+        default="kaiming_init",
+        title="LoRA A matrix initializer",
+        description="Custom initialization function for LoRA A matrix. Default to Kaiming uniform initialization.",
+    )
+
+    lora_b_init: str = ModeloptField(
+        default="zero_init",
+        title="LoRA B matrix initializer",
+        description="Custom initialization function for LoRA B matrix. Default to zero initialization.",
+    )
+
+    @field_validator("lora_a_init", "lora_b_init")
+    @classmethod
+    def validate_init_method(cls, v):
+        """Validate initialization method is supported."""
+        valid_methods = {"kaiming_init", "zero_init"}
+        if v not in valid_methods:
+            raise ValueError(
+                f"Invalid initialization method: {v}. Supported methods: {', '.join(valid_methods)}"
+            )
+        return v
+
+    @field_validator("rank")
+    @classmethod
+    def validate_rank(cls, v):
+        """Validate rank is positive."""
+        if v < 1:
+            raise ValueError("rank must be a positive integer")
+        return v
+
+    @field_validator("scale")
+    @classmethod
+    def validate_scale(cls, v):
+        """Validate scale is positive."""
+        if v <= 0:
+            raise ValueError("scale must be a positive number")
+        return v
+
+
+# Type alias for adapter configuration
+PEFTAdapterCfgType = dict[str | Callable, PEFTAttributeConfig | dict]
+
+
+class PEFTConfig(ModeloptBaseConfig):
+    """Default configuration for ``peft`` mode."""
+
+    adapter_name: str = ModeloptField(
+        default="default",
+        title="Adapter name",
+        description="Name of the adapter to create or update.",
+        validate_default=True,
+    )
+
+    adapter_cfg: PEFTAdapterCfgType = ModeloptField(
+        default={"*": {"rank": 64}},
+        title="Adapter configuration",
+        description="Configuration for adapters. Maps module patterns to PEFTAttributeConfig or dict.",
+        validate_default=True,
+    )
+
+    adapter_type: str = ModeloptField(
+        default="lora",
+        title="Adapter type",
+        description="Type of PEFT adapter to use. Currently only 'lora' is supported.",
+        validate_default=True,
+    )
+
+    freeze_base_model: bool = ModeloptField(
+        default=True,
+        title="Freeze base weights during training",
+        description="Whether to freeze the base model weights; in most cases, this should be set to True.",
+        validate_default=True,
+    )
+
+    freeze_lora_weights: bool = ModeloptField(
+        default=False,
+        title="Freeze lora weights during training",
+        description="Whether to freeze the lora model weights; in most cases, this should be set to False.",
+        validate_default=True,
+    )
+
+    @field_validator("adapter_type")
+    @classmethod
+    def validate_adapter_type(cls, v):
+        """Validate adapter type."""
+        if v not in ["lora"]:
+            raise ValueError(f"Unsupported adapter type: {v}. Only 'lora' is currently supported.")
+        return v
+
+    @field_validator("adapter_cfg")
+    @classmethod
+    def validate_adapter_cfg(cls, v):
+        """Validate and convert adapter configurations."""
+        validated_cfg = {}
+        for key, value in v.items():
+            if isinstance(value, dict) and not isinstance(value, PEFTAttributeConfig):
+                # Convert dict to PEFTAttributeConfig to trigger validation
+                try:
+                    validated_cfg[key] = PEFTAttributeConfig(**value)
+                except Exception as e:
+                    raise ValueError(f"Invalid adapter configuration for '{key}': {e}")
+            else:
+                validated_cfg[key] = value
+        return validated_cfg
+
+
+class ExportPEFTConfig(ModeloptBaseConfig):
+    """An empty config."""