Merge pull request #1 from parmanu-lcs2/parmanu-lcs2/main

Implemented MonteCLoRA into LoRA variants

Author: iamsecretlyflash

.gitignore

@@ -7,6 +7,7 @@ __pycache__/
 *.so
 
 # Distribution / packaging
+test_imp/
 .Python
 build/
 develop-eggs/

src/peft/__init__.py

@@ -25,6 +25,7 @@
     AutoPeftModelForTokenClassification,
 )
 from .config import PeftConfig, PromptLearningConfig
+from .helpers import MonteCLoRATrainerMixin
 from .mapping import (
     PEFT_TYPE_TO_CONFIG_MAPPING,
     PEFT_TYPE_TO_MIXED_MODEL_MAPPING,
@@ -87,7 +88,6 @@
     MissConfig,
     MissModel,
     MonteCLoraConfig,
-    MonteCLoraModel,
     MultitaskPromptTuningConfig,
     MultitaskPromptTuningInit,
     OFTConfig,
@@ -202,8 +202,8 @@
     "LoraRuntimeConfig",
     "MissConfig",
     "MissModel",
+    "MonteCLoRATrainerMixin",
     "MonteCLoraConfig",
-    "MonteCLoraModel",
     "MultitaskPromptTuningConfig",
     "MultitaskPromptTuningInit",
     "OFTConfig",

src/peft/helpers.py

@@ -251,6 +251,93 @@ def disable_input_dtype_casting(model: nn.Module, active: bool = True):
                 module.cast_input_dtype_enabled = original_values[name]
 
 
+class MonteCLoRATrainerMixin:
+    """
+    Mixin class for adding MonteCLoRA variational loss to the Trainer's compute_loss method.
+
+    This mixin can be used with any Trainer class (e.g., Trainer, SFTTrainer) to add support for
+    MonteCLoRA's variational regularization during training.
+
+    Example:
+        ```python
+        from transformers import Trainer
+        from peft import get_peft_model, LoraConfig
+        from peft.helpers import MonteCLoRATrainerMixin
+        from peft.tuners.monteclora_new import MonteCLoraConfig
+
+        #custom trainer that supports MonteCLoRA
+        class MonteCLoRATrainer(MonteCLoRATrainerMixin, Trainer):
+            pass
+
+        # Configure LoRA with MonteCLoRA
+        monteclora_config = MonteCLoraConfig(
+            monteclora_n=8,
+            sample_scaler=1e-4,
+            kl_loss_weight=1e-5,
+        )
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["q_proj", "v_proj"],
+            use_monteclora=True,
+            monteclora_config=monteclora_config,
+        )
+
+        # Get PEFT model and train
+        model = get_peft_model(base_model, lora_config)
+        trainer = MonteCLoRATrainer(model=model, args=training_args, ...)
+        trainer.train()
+        ```
+    """
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        """
+        Compute loss with MonteCLoRA variational regularization.
+
+        This method extends the standard compute_loss by adding the variational loss
+        (KL divergence + entropy) from MonteCLoRA samplers to the task loss.
+
+        Args:
+            model: The model being trained
+            inputs: Input batch
+            return_outputs: Whether to return model outputs along with loss
+            **kwargs: Additional arguments
+
+        Returns:
+            loss or (loss, outputs) depending on return_outputs
+        """
+        # 1. Compute the standard task loss
+        if return_outputs:
+            task_loss, outputs = super().compute_loss(model, inputs, return_outputs=True, **kwargs)
+        else:
+            task_loss = super().compute_loss(model, inputs, return_outputs=False, **kwargs)
+            outputs = None
+
+        # 2. Calculate Variational Loss (KLD + Entropy) from MonteCLoRA samplers
+        var_loss_sum = 0.0
+        num_monte_layers = 0
+
+        # Iterate through modules to find MonteCLoRA samplers
+        for name, module in model.named_modules():
+            # Check if this is a MonteCLoRASampler by checking for the get_variational_loss method
+            if hasattr(module, "get_variational_loss") and module.__class__.__name__ == "MonteCLoRASampler":
+                try:
+                    kl_loss, entropy_loss = module.get_variational_loss()
+                    var_loss_sum += kl_loss + entropy_loss
+                    num_monte_layers += 1
+                except Exception:
+                    # Silently ignore if get_variational_loss fails
+                    pass
+
+        # 3. Normalize the Variational Loss
+        regularization_loss = 0.0
+        if num_monte_layers > 0:
+            regularization_loss = var_loss_sum / num_monte_layers
+
+        # 4. Combine losses
+        total_loss = task_loss + regularization_loss
+
+        return (total_loss, outputs) if return_outputs else total_loss
 class DoraCaching:
     """Context manager to enable DoRA caching, which improves speed of DoRA inference at the expense of memory.
 

src/peft/tuners/__init__.py

@@ -45,7 +45,7 @@
 )
 from .miss import MissConfig, MissModel
 from .mixed import MixedModel
-from .monteclora import MonteCLoraConfig, MonteCLoraModel
+from .monteclora import MonteCLoraConfig
 from .multitask_prompt_tuning import MultitaskPromptEmbedding, MultitaskPromptTuningConfig, MultitaskPromptTuningInit
 from .oft import OFTConfig, OFTModel
 from .osf import OSFConfig, OSFModel
@@ -106,7 +106,6 @@
     "MissModel",
     "MixedModel",
     "MonteCLoraConfig",
-    "MonteCLoraModel",
     "MultitaskPromptEmbedding",
     "MultitaskPromptTuningConfig",
     "MultitaskPromptTuningInit",

src/peft/tuners/lora/config.py

@@ -693,6 +693,26 @@ class LoraConfig(PeftConfig):
             )
         },
     )
+    use_monteclora: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Enable MonteCLoRA (Monte Carlo Low-Rank Adaptation). This technique introduces variational "
+                "inference into LoRA by adding Monte Carlo sampling to the adapter weights during training. "
+                "This can improve model performance and uncertainty estimation. When enabled, you should also "
+                "provide `monteclora_config` with the MonteCLoRA hyperparameters."
+            )
+        },
+    )
+    monteclora_config: Optional[MonteCLoraConfig] = field(  # noqa: F821
+        default=None,
+        metadata={
+            "help": (
+                "The configuration of MonteCLoRA. If this is passed along with `use_monteclora=True`, then "
+                "MonteCLoRA will be used to add variational sampling to the LoRA adapters."
+            )
+        },
+    )
     # Enables replicating layers in a model to expand it to a larger model.
     layer_replication: Optional[list[tuple[int, int]]] = field(
         default=None,
@@ -832,6 +852,20 @@ def __post_init__(self):
         elif self.init_lora_weights != "corda" and self.corda_config is not None:
             warnings.warn("`corda_config` specified but will be ignored when `init_lora_weights` is not 'corda'.")
 
+        # Handle MonteCLoRA configuration
+        if self.use_monteclora:
+            from peft.tuners.monteclora.config import MonteCLoraConfig
+
+            if self.monteclora_config is None:
+                warnings.warn(
+                    "`use_monteclora=True` but `monteclora_config` is not specified. Using default MonteCLoRA config."
+                )
+                self.monteclora_config = MonteCLoraConfig()
+            elif isinstance(self.monteclora_config, dict):
+                self.monteclora_config = MonteCLoraConfig(**self.monteclora_config)
+        elif self.monteclora_config is not None:
+            warnings.warn("`monteclora_config` specified but will be ignored when `use_monteclora=False`.")
+
         if self.lora_bias:
             if self.init_lora_weights not in (True, False):
                 raise ValueError(

src/peft/tuners/lora/layer.py

@@ -147,6 +147,21 @@ def resolve_lora_variant(self, *, config: LoraConfig, **kwargs) -> Optional[Lora
 
     def update_layer(
         self,
+        adapter_name,
+        r,
+        lora_alpha,
+        lora_dropout,
+        init_lora_weights,
+        use_rslora,
+        use_dora: bool = False,
+        use_alora: bool = False,
+        use_qalora: bool = False,
+        use_monteclora: bool = False,
+        lora_bias: bool = False,
+        arrow_config: ArrowConfig = None,
+        monteclora_config=None,
+        qalora_group_size: int = 32,
+        inference_mode: bool = False,
         adapter_name: str,
         r: int,
         lora_alpha: int,
@@ -174,7 +189,15 @@ def update_layer(
                 PeftWarning,
             )
 
-        lora_variant = self.resolve_lora_variant(config=config)
+        lora_variant = self.resolve_lora_variant(
+            use_dora=use_dora,
+            use_alora=use_alora,
+            use_qalora=use_qalora,
+            use_monteclora=use_monteclora,
+            qalora_group_size=qalora_group_size,
+            arrow_config=arrow_config,
+            monteclora_config=monteclora_config,
+        )
         if lora_variant is not None:
             self.lora_variant[adapter_name] = lora_variant
 
@@ -732,6 +755,14 @@ def __init__(
         r: int = 0,
         lora_alpha: int = 1,
         is_target_conv_1d_layer: bool = False,
+        init_lora_weights: Union[bool, str] = True,
+        use_rslora: bool = False,
+        use_dora: bool = False,
+        use_alora: bool = False,
+        use_monteclora: bool = False,
+        arrow_config: ArrowConfig = None,
+        monteclora_config=None,
+        lora_bias: bool = False,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -743,24 +774,32 @@ def __init__(
             adapter_name,
             r,
             lora_alpha=lora_alpha,
-            config=config,
-            **kwargs,
+            lora_dropout=lora_dropout,
+            init_lora_weights=init_lora_weights,
+            use_rslora=use_rslora,
+            use_dora=use_dora,
+            use_alora=use_alora,
+            use_monteclora=use_monteclora,
+            lora_bias=lora_bias,
+            arrow_config=arrow_config,
+            monteclora_config=monteclora_config,
         )
         self.is_target_conv_1d_layer = is_target_conv_1d_layer
 
-    def resolve_lora_variant(self, config: LoraConfig, **kwargs) -> Optional[LoraVariant]:
-        if config.arrow_config is not None:
+    def resolve_lora_variant(
+        self, *, arrow_config: ArrowConfig, use_dora: bool, use_alora: bool, use_monteclora: bool = False, **kwargs
+    ) -> Optional[LoraVariant]:
+        if arrow_config is not None:
             from .variants import ArrowLinearVariant
 
             return ArrowLinearVariant()
 
-        if config.use_bdlora is not None:
-            from .variants import BdLoraLinearVariant
+        if use_monteclora:
+            from peft.tuners.monteclora.variant import MonteCLoraLinearVariant
 
-            return BdLoraLinearVariant()
+            return MonteCLoraLinearVariant()
 
-        use_alora = config.alora_invocation_tokens is not None
-        if not config.use_dora and not use_alora:
+        if not use_dora and not use_alora:
             return None
 
         from .variants import ALoraLinearVariant, DoraLinearVariant

src/peft/tuners/lora/model.py

@@ -206,7 +206,19 @@ def _create_and_replace(
         kwargs = {
             "r": r,
             "lora_alpha": alpha,
-            "target_name": current_key,
+            "lora_dropout": lora_config.lora_dropout,
+            "fan_in_fan_out": lora_config.fan_in_fan_out,
+            "init_lora_weights": lora_config.init_lora_weights,
+            "use_rslora": lora_config.use_rslora,
+            "use_dora": lora_config.use_dora,
+            "use_alora": lora_config.alora_invocation_tokens is not None,
+            "use_qalora": lora_config.use_qalora,
+            "use_monteclora": lora_config.use_monteclora,
+            "qalora_group_size": lora_config.qalora_group_size,
+            "monteclora_config": lora_config.monteclora_config,
+            "ephemeral_gpu_offload": lora_config.runtime_config.ephemeral_gpu_offload,
+            "lora_bias": lora_config.lora_bias,
+            "arrow_config": lora_config.arrow_config,
             "loaded_in_8bit": getattr(self.model, "is_loaded_in_8bit", False),
             "loaded_in_4bit": getattr(self.model, "is_loaded_in_4bit", False),
             "parameter_name": parameter_name,
@@ -236,8 +248,15 @@ def _create_and_replace(
                 adapter_name,
                 r,
                 lora_alpha=alpha,
-                target_name=current_key,
-                config=lora_config,
+                lora_dropout=lora_config.lora_dropout,
+                init_lora_weights=lora_config.init_lora_weights,
+                use_rslora=lora_config.use_rslora,
+                use_dora=lora_config.use_dora,
+                use_monteclora=lora_config.use_monteclora,
+                lora_bias=lora_config.lora_bias,
+                arrow_config=lora_config.arrow_config,
+                monteclora_config=lora_config.monteclora_config,
+                inference_mode=lora_config.inference_mode,
             )
         else:
             if isinstance(target, ParamWrapper) and (parameter_name == target.parameter_name):

src/peft/tuners/monteclora/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023-present the HuggingFace Inc. team.
+# Copyright 2026-present the HuggingFace Inc. team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,35 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from peft.import_utils import is_bnb_4bit_available, is_bnb_available, is_eetq_available
-from peft.utils import register_peft_method
-
 from .config import MonteCLoraConfig
-from .layer import MonteCLoraLinear
-from .model import MonteCLoraModel
-
-
-__all__ = ["MonteCLoraConfig", "MonteCLoraLinear", "MonteCLoraModel"]
-
-register_peft_method(
-    name="monteclora", prefix="lora_", config_cls=MonteCLoraConfig, model_cls=MonteCLoraModel, is_mixed_compatible=True
-)
-
-
-def __getattr__(name):
-    if (name == "Linear8bitLt") and is_bnb_available():
-        from peft.tuners.lora.bnb import Linear8bitLt
-
-        return Linear8bitLt
-
-    if (name == "Linear4bit") and is_bnb_4bit_available():
-        from peft.tuners.lora.bnb import Linear4bit
-
-        return Linear4bit
-
-    if (name == "EetqLoraLinear") and is_eetq_available():
-        from peft.tuners.lora.eetq import EetqLoraLinear
+from .variant import MonteCLoraLinearVariant
 
-        return EetqLoraLinear
 
-    raise AttributeError(f"module {__name__} has no attribute {name}")
+__all__ = ["MonteCLoraConfig", "MonteCLoraLinearVariant"]