Updated LoraConfig to subclass from peft.LoraConfig

Author: romitjain

docs/tuning-techniques.md

@@ -24,7 +24,7 @@
 
 ## LoRA Tuning Example
 
-Set `peft_method` to `"lora"`. You can additionally pass any arguments from [LoraConfig](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/tuning/config/peft_config.py#L21).
+Set `peft_method` to `"lora"`. You can additionally pass any arguments from [LoraConfig](https://huggingface.co/docs/peft/en/package_reference/lora#peft.LoraConfig).
 ```py
 # Args you can pass
 r: int =8 
@@ -340,7 +340,7 @@ You can see details on a sample configuration of Accelerated GPTQ-LoRA [here](ht
 
 To use GPTQ-LoRA technique, you can set the `quantized_lora_config` defined [here](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/tuning/config/acceleration_configs/quantized_lora_config.py). See the Notes section of FMS Acceleration doc [below](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/README.md#fms-acceleration) for usage. The only kernel we are supporting currently is `triton_v2`.
 
-In addition, LoRA tuning technique is required to be used, set `peft_method` to `"lora"` and pass any arguments from [LoraConfig](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/tuning/config/peft_config.py#L21).
+In addition, LoRA tuning technique is required to be used, set `peft_method` to `"lora"` and pass any arguments from [LoraConfig](https://huggingface.co/docs/peft/en/package_reference/lora#peft.LoraConfig).
 
 Example command to run:
 

tuning/config/peft_config.py

@@ -18,6 +18,7 @@
 from typing import List
 
 # Third Party
+from peft import LoraConfig as _LoraConfig
 from transformers.utils.quantization_config import Mxfp4Config as HfMxfp4Config
 
 
@@ -40,50 +41,26 @@ def to_hf_config(self):
 
 
 @dataclass
-class LoraConfig:
+class LoraConfig(_LoraConfig):
     """
-    This is the configuration class to store the configuration of a [`LoraModel`].
+    This is the configuration class that extends peft.LoraConfig with a few defaults.
 
     Args:
-        r (`int`):
-            Lora attention dimension (the "rank").
-        target_modules (List[str]]):
-            The names of the modules to apply the adapter to. \
-            If this is specified, only the modules with the specified \
-            names will be replaced. Please specify modules as per model architecture. \
-            If the value is ["all-linear"], \
-            then LORA selects all linear and Conv1D modules as per model architecture, \
-            except for the output layer.
         lora_alpha (`int`):
             The alpha parameter for Lora scaling.
         lora_dropout (`float`):
             The dropout probability for Lora layers.
-        bias (`str`):
-            Bias type for LoRA. Can be 'none', 'all' or 'lora_only'. \
-            If 'all' or 'lora_only', the corresponding biases will be updated during training. \
-            Be aware that this means that, even when disabling the adapters, the model \
-            will not produce the same output as the base model would have without adaptation.
     """
-
-    r: int = 8
     lora_alpha: int = 32
-    target_modules: List[str] = field(
-        default=None,
-        metadata={
-            "help": "The names of the modules to apply LORA to. LORA selects modules which either \
-            completely match or "
-            'end with one of the strings. If the value is ["all-linear"], \
-            then LORA selects all linear and Conv1D '
-            "modules except for the output layer."
-        },
-    )
-    target_parameters: List[str] = field(
-        default=None,
-        metadata={"help": "The names/regex of the parameters to apply LORA to"},
-    )
-    bias = "none"
     lora_dropout: float = 0.05
 
+    def __post_init__(self):
+        # If target_modules is a single-element list, convert it into a plain string
+        if self.target_modules == ["all-linear"]:
+            self.target_modules = "all-linear"
+
+        super().__post_init__()
+
 
 @dataclass
 class PromptTuningConfig:

tuning/sft_trainer.py

@@ -71,7 +71,7 @@ def train(
     data_args: configs.DataArguments,
     train_args: configs.TrainingArguments,
     peft_config: Optional[  # pylint: disable=redefined-outer-name
-        Union[peft_config.LoraConfig, LoraConfig, peft_config.PromptTuningConfig]
+        Union[LoraConfig, peft_config.PromptTuningConfig]
     ] = None,
     quantization_config: Optional[peft_config.Mxfp4Config] = None,
     trainer_controller_args: TrainerControllerCallback = None,
@@ -92,8 +92,7 @@ def train(
         model_args: tuning.config.configs.ModelArguments
         data_args: tuning.config.configs.DataArguments
         train_args: tuning.config.configs.TrainingArguments
-        peft_config: peft_config.LoraConfig for Lora tuning | \
-        LoraConfig (peft.LoraConfig): for activated Lora (aLoRA) tuning | \
+        peft_config: LoraConfig (peft.LoraConfig): for activated Lora (aLoRA) tuning | \
         peft_config.PromptTuningConfig for prompt tuning | \
         None for full fine tuning
             The peft configuration to pass to trainer
@@ -110,7 +109,7 @@ def train(
                               tracker with automatically be added.
         exp_metadata: Dict of key value pairs passed to train to be recoreded by the tracker.
         quantized_lora_config: tuning.config.acceleration_configs.QuantizedLoraConfig \
-            Should be used in combination with peft_config.LoraConfig for Lora tuning \
+            Should be used in combination with LoraConfig for Lora tuning \
         fusedops_kernels_config: tuning.config.acceleration_configs.FusedOpsAndKernelsConfig \
             Should be used in combination with quantized_lora_config. Also currently 
             fused_lora and fast_kernels must used together (may change in future). \
@@ -855,7 +854,7 @@ def main():
             sys.exit(INTERNAL_ERROR_EXIT_CODE)
 
     if isinstance(
-        tune_config, (peft_config.LoraConfig, LoraConfig)
+        tune_config, LoraConfig
     ):  # aLoraConfig subclasses LoraConfig
         try:
             if training_args.save_model_dir:

tuning/utils/config_utils.py

@@ -113,9 +113,10 @@ def get_hf_peft_config(task_type, tuning_config, tokenizer_name_or_path):
         hf_peft_config = alora_config
     elif isinstance(tuning_config, peft_config.LoraConfig):
         lora_config = asdict(tuning_config)
-        if lora_config["target_modules"] == ["all-linear"]:
-            lora_config["target_modules"] = "all-linear"
-        hf_peft_config = HFLoraConfig(task_type=task_type, **lora_config)
+
+        if not hasattr(lora_config, "task_type"):
+            lora_config["task_type"]=task_type
+        hf_peft_config = HFLoraConfig(**lora_config)
     elif isinstance(tuning_config, peft_config.PromptTuningConfig):
         hf_peft_config = HFPromptTuningConfig(
             task_type=task_type,