Add Validation for Invalid task_type in PEFT Configurations

[d-kleine]

Fixes #2203

This PR adds validation for the task_type parameter across all PEFT types (e.g., LoRA, Prefix Tuning, etc.) to ensure that only valid task types, as defined in the TaskType enum, are used. This resolves the issue where invalid or non-sense task types (e.g., "BLABLA") or typos (e.g., "CUASAL_LM" instead of "CAUSAL_LM") could be passed without raising an error, potentially causing unexpected behavior.

Why This Change Is Necessary:

• Prevents Silent Failures: Without this validation, users could pass incorrect or unsupported task types without receiving any feedback. This could lead to silent failures or unexpected behavior during training or inference.
• Consistency: Ensures that all PEFT types adhere to the same validation rules for task_type, maintaining consistency across different PEFT configurations.

Tests:

• Ensured that invalid task types raise appropriate errors.
• Verified that valid task types continue to work as expected.

Example Scenario:

Before this PR, the following code would not raise an error despite using an invalid task type:

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

# Step 1: Load the base model and tokenizer
model_name_or_path = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# Step 2: Define the LoRA configuration
lora_config = LoraConfig(
    r=16,  
    lora_alpha=32, 
    lora_dropout=0.1,  
    target_modules=["q_proj", "v_proj"],  
    bias="none",  
    task_type="BLABLA"  # Invalid task type 
)

# Step 3: Wrap the base model with LoRA using the configuration
peft_model = get_peft_model(model, lora_config)

With this PR, attempting to use an invalid task type like "BLABLA" will now raise a ValueError, ensuring that only valid task types are accepted:

ValueError: Invalid task type: 'BLABLA'. Must be one of the following task types: ['SEQ_CLS', 'SEQ_2_SEQ_LM', 'CAUSAL_LM', 'TOKEN_CLS', 'QUESTION_ANS', 'FEATURE_EXTRACTION'].

For multiple PEFT methods:

INVALID task_type

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import (
    LoraConfig, 
    PromptTuningConfig, 
    PrefixTuningConfig,
    get_peft_model
)

# Step 1: Load the base model and tokenizer
model_name_or_path = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# Step 2: Define configurations for each PEFT method with an invalid task type
configs = [
    LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.1,
        target_modules=["q_proj", "v_proj"],
        bias="none",
        task_type="BLABLA"  # Invalid task type
    ),
    PromptTuningConfig(
        num_virtual_tokens=20,
        task_type="BLABLA"  # Invalid task type
    ),
    PrefixTuningConfig(
        num_virtual_tokens=30,
        task_type="BLABLA"  # Invalid task type
    ),
    # Add more configurations for other PeftTypes if necessary...
]

# Step 3: Test each configuration and check if it raises a ValueError for invalid task type
for config in configs:
    try:
        print(f"Testing {config.__class__.__name__}...")
        peft_model = get_peft_model(model, config)
        print("OK")
    except ValueError as e:
        print(f"Error: {e}")

Testing LoraConfig...
Error: Invalid task type: 'BLABLA'. Must be one of the following task types: ['SEQ_CLS', 'SEQ_2_SEQ_LM', 'CAUSAL_LM', 'TOKEN_CLS', 'QUESTION_ANS', 'FEATURE_EXTRACTION'].
Testing PromptTuningConfig...
Error: Invalid task type: 'BLABLA'. Must be one of the following task types: ['SEQ_CLS', 'SEQ_2_SEQ_LM', 'CAUSAL_LM', 'TOKEN_CLS', 'QUESTION_ANS', 'FEATURE_EXTRACTION'].
Testing PrefixTuningConfig...
Error: Invalid task type: 'BLABLA'. Must be one of the following task types: ['SEQ_CLS', 'SEQ_2_SEQ_LM', 'CAUSAL_LM', 'TOKEN_CLS', 'QUESTION_ANS', 'FEATURE_EXTRACTION'].

VALID task_type

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import (
    LoraConfig, 
    PromptTuningConfig, 
    PrefixTuningConfig,
    get_peft_model
)

# Step 1: Load the base model and tokenizer
model_name_or_path = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# Step 2: Define configurations for each PEFT method with an valid task type
configs = [
    LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.1,
        target_modules=["q_proj", "v_proj"],
        bias="none",
        task_type="CAUSAL_LM"  # Valid task type
    ),
    PromptTuningConfig(
        num_virtual_tokens=20,
        task_type="CAUSAL_LM"  # Valid task type
    ),
    PrefixTuningConfig(
        num_virtual_tokens=30,
        task_type="CAUSAL_LM"  # Valid task type
    ),
]

# Step 3: Test each configuration and check if it raises a ValueError for invalid task type
for config in configs:
    try:
        print(f"Testing {config.__class__.__name__}...")
        peft_model = get_peft_model(model, config)
        print("OK")
    except ValueError as e:
        print(f"Error: {e}")

Testing LoraConfig...
OK
Testing PromptTuningConfig...
OK
Testing PrefixTuningConfig...
OK

[d-kleine]

@BenjaminBossan What do you think of this lean implementation?

[BenjaminBossan]

@d-kleine Thanks for the PR. Right now, I'm at a company offsite and don't have the means to do proper reviews. I'll probably review at the start of next week.

[BenjaminBossan]

Thanks for the PR. Looking at this implementation, I think we need to adjust the logic a bit. Right now, if the user does not explicitly pass task_type, it will be None. Since None is not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING, users may now get an error even though their code might work just fine.

I think the best way to check this is instead: If task_type is not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING and if task_type is not None, we know that the user provided a value but that value is misspelled. Let's only raise an error in that case.

I would also suggest to check this in a different place. Right now, this is checked inside of get_peft_model. But how about checking it already when the config is created? That way, we can give the feedback as early as possible. I would thus suggest to move the check to peft.config.PeftConfig.__post_init__. In addition to that, I think we need to ensure that all configs that inherit, i.e. LoraConfig, PrefixTuningConfig, etc. all call super().__post_init__() in their own __post_init__ methods. It's a bit annoying to have to change this in each config, but I think it's worth it.

Let me know your thoughts overall.

[d-kleine]

I see what you mean. I have implemented the changes, they work good so far, returning an error when setting up the config with either no task_type is passed (None) or the one passed does not match with the task types defined in the mappings.

About the super().__post_init__(), I agree this also a good idea for this project in general. But I have noticed most parents (e.g. PeftConfig, PromptLearningConfig) don't have a __post_init__() method yet, thus calling super().__post_init__() in the def __post_init__(self): in would result in an error message then

parent:

@dataclass
class PeftConfig(PeftConfigMixin):
    ... # no `def __post_init__(self):` here

child:

@dataclass
class LoraConfig(PeftConfig):
...
    def __post_init__(self):
        super().__post_init__() # error as not defined in parent

So, what to do? Something like this? ⬇️

def __post_init__(self):
    if hasattr(super(), '__post_init__'):
        super().__post_init__()

[BenjaminBossan]

Hmm, I think this check for the presence of __post_init__ should not be necessary. I tried this locally:

# peft/config.py in PeftConfigMixin
    def __post_init__(self):
        if (self.task_type is not None) and (self.task_type not in list(TaskType)):
            raise ValueError(f"Invalid task type: '{self.task_type}'. Must be one of the following task types: {', '.join(TaskType)}.")

# peft/tuners/lora/config.py
    def __post_init__(self):
        super().__post_init__()

        ...  # rest of the code

Then running:

from peft import LoraConfig

lora_config_1 = LoraConfig(task_type=None)
lora_config_1 = LoraConfig(task_type="CAUSAL_LM")
lora_config_2 = LoraConfig(task_type="CAUSAL_LMX")

I get the expected error on the last config:

ValueError: Invalid task type: 'CAUSAL_LMX'. Must be one of the following task types: SEQ_CLS, SEQ_2_SEQ_LM, CAUSAL_LM, TOKEN_CLS, QUESTION_ANS, FEATURE_EXTRACTION.