Fix CI dev test TypeError: unexpected keyword argument 'load_in_4bit'

docs/source/multi_adapter_rl.md

@@ -90,7 +90,7 @@ model = AutoModelForCausalLMWithValueHead.from_pretrained(
     model_name,
     peft_config=lora_config,
     reward_adapter=rm_adapter_id,
-    load_in_8bit=True,
+    quantization_config=BitsAndBytesConfig(load_in_8bit=True),
 )
 
 ...

tests/test_dpo_trainer.py

@@ -642,6 +642,7 @@ def test_dpo_lora_save(self):
     def test_dpo_lora_bf16_autocast_llama(self):
         # Note this test only works on compute capability > 7 GPU devices
         from peft import LoraConfig
+        from transformers import BitsAndBytesConfig
 
         model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
         tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -655,7 +656,9 @@ def test_dpo_lora_bf16_autocast_llama(self):
         )
 
         # lora model
-        model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id, quantization_config=BitsAndBytesConfig(load_in_4bit=True)
+        )
 
         training_args = DPOConfig(
             output_dir=self.tmp_dir,
@@ -725,6 +728,7 @@ def test_dpo_lora_bf16_autocast_llama(self):
     )
     def test_dpo_lora_bf16_autocast(self, loss_type, pre_compute, gen_during_eval):
         from peft import LoraConfig
+        from transformers import BitsAndBytesConfig
 
         lora_config = LoraConfig(
             r=16,
@@ -735,7 +739,9 @@ def test_dpo_lora_bf16_autocast(self, loss_type, pre_compute, gen_during_eval):
         )
 
         # lora model
-        model = AutoModelForCausalLM.from_pretrained(self.model_id, load_in_4bit=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            self.model_id, quantization_config=BitsAndBytesConfig(load_in_4bit=True)
+        )
 
         training_args = DPOConfig(
             output_dir=self.tmp_dir,

tests/test_peft_models.py

@@ -101,17 +101,20 @@ def test_create_bnb_peft_model_from_config(self):
         Simply creates a peft model and checks that it can be loaded.
         """
         from bitsandbytes.nn import Linear8bitLt
+        from transformers import BitsAndBytesConfig
 
         trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(
-            self.causal_lm_model_id, peft_config=self.lora_config, load_in_8bit=True
+            self.causal_lm_model_id,
+            peft_config=self.lora_config,
+            quantization_config=BitsAndBytesConfig(load_in_8bit=True),
         )
         # Check that the number of trainable parameters is correct
         nb_trainable_params = sum(p.numel() for p in trl_model.parameters() if p.requires_grad)
         assert nb_trainable_params == 905
         assert isinstance(trl_model.pretrained_model.model.model.layers[0].mlp.gate_proj, Linear8bitLt)
 
         causal_lm_model = AutoModelForCausalLM.from_pretrained(
-            self.causal_lm_model_id, load_in_8bit=True, device_map="auto"
+            self.causal_lm_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True), device_map="auto"
         )
         trl_model = AutoModelForCausalLMWithValueHead.from_pretrained(causal_lm_model, peft_config=self.lora_config)
         # Check that the number of trainable parameters is correct

trl/models/modeling_base.py

@@ -149,8 +149,13 @@ class and the arguments that are specific to trl models. The kwargs also support
 
         current_device = cls._get_current_device()
         if isinstance(pretrained_model_name_or_path, str):
-            is_loaded_in_8bit = pretrained_kwargs["load_in_8bit"] if "load_in_8bit" in pretrained_kwargs else False
-            is_loaded_in_4bit = pretrained_kwargs["load_in_4bit"] if "load_in_4bit" in pretrained_kwargs else False
+            quantization_config = pretrained_kwargs.get("quantization_config", None)
+            if quantization_config is not None:
+                is_loaded_in_8bit = getattr(quantization_config, "load_in_8bit", False)
+                is_loaded_in_4bit = getattr(quantization_config, "load_in_4bit", False)
+            else:
+                is_loaded_in_8bit = pretrained_kwargs["load_in_8bit"] if "load_in_8bit" in pretrained_kwargs else False
+                is_loaded_in_4bit = pretrained_kwargs["load_in_4bit"] if "load_in_4bit" in pretrained_kwargs else False
         else:
             is_loaded_in_8bit = getattr(pretrained_model_name_or_path, "is_loaded_in_8bit", False)
             is_loaded_in_4bit = getattr(pretrained_model_name_or_path, "is_loaded_in_4bit", False)