diff --git a/tests/test_reward_trainer.py b/tests/test_reward_trainer.py index 51ea1183f7..5a858b2f38 100644 --- a/tests/test_reward_trainer.py +++ b/tests/test_reward_trainer.py @@ -189,6 +189,28 @@ def test_train_lora_pretokenized(self): new_param = trainer.model.get_parameter(n) self.assertTrue(torch.allclose(param, new_param, atol=1e-12, rtol=1e-12)) + @require_peft + def test_all_linear_user_warning(self): + peft_config = LoraConfig( + task_type=TaskType.SEQ_CLS, + inference_mode=False, + r=8, + lora_alpha=32, + lora_dropout=0.1, + target_modules="all-linear", + ) + with tempfile.TemporaryDirectory() as tmp_dir: + dummy_dataset = load_dataset("trl-internal-testing/zen", "conversational_preference", split="train") + training_args = RewardConfig(output_dir=tmp_dir, max_steps=3, report_to="none") + with self.assertWarns(UserWarning, msg="You passed target_modules="): + _ = RewardTrainer( + model=self.model, + args=training_args, + processing_class=self.tokenizer, + train_dataset=dummy_dataset, + peft_config=peft_config, + ) + def test_margin(self): with tempfile.TemporaryDirectory() as tmp_dir: dummy_dataset_dict = { diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py index 063fe5e8e8..6069104608 100644 --- a/trl/trainer/reward_trainer.py +++ b/trl/trainer/reward_trainer.py @@ -25,6 +25,7 @@ from accelerate import PartialState from accelerate.utils import gather_object from datasets import Dataset +from peft.utils import INCLUDE_LINEAR_LAYERS_SHORTHAND from transformers import ( BaseImageProcessor, DataCollator, @@ -159,6 +160,19 @@ def __init__( model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + # Warn if the user passes "all-linear" for the target_modules + target_modules = ( + peft_config.get("target_modules", None) + if isinstance(peft_config, dict) + else peft_config.target_modules + ) + if target_modules == INCLUDE_LINEAR_LAYERS_SHORTHAND: + warnings.warn( + f"You passed target_modules='{INCLUDE_LINEAR_LAYERS_SHORTHAND}' in the peft_config. " + "This will result in all linear layers except the output layer being adapted. " + "This could negatively impact the performance of the reward model as the output layer (used for scoring of chosen and rejected completions) will not be adapted or trained.", + UserWarning, + ) model = get_peft_model(model, peft_config) # Disable dropout in the model