[CI] Modify tests to handle device allocation for models (#3962)

Author: kashif

tests/test_modeling_geometric_mixture_wrapper.py

@@ -25,16 +25,17 @@ class TestGeometricMixtureWrapper(TrlTestCase):
     def setUp(self):
         super().setUp()
         model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
-        self.model = AutoModelForCausalLM.from_pretrained(model_id)
-        self.ref_model = create_reference_model(self.model)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = AutoModelForCausalLM.from_pretrained(model_id).to(self.device)
+        self.ref_model = create_reference_model(self.model).to(self.device)
         self.generation_config = GenerationConfig.from_pretrained(model_id)
         self.mixture_coef = 0.5
         self.wrapper = GeometricMixtureWrapper(
             self.model, self.ref_model, self.generation_config, mixture_coef=self.mixture_coef
         )
 
     def test_forward(self):
-        input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        input_ids = torch.tensor([[1, 2, 3, 4, 5]], device=self.device)
         attention_mask = torch.ones_like(input_ids)
 
         output = self.wrapper(input_ids=input_ids, attention_mask=attention_mask)
@@ -44,7 +45,7 @@ def test_forward(self):
         self.assertEqual(output.logits.shape, (1, 5, self.model.config.vocab_size))
 
     def test_mixture_coefficient(self):
-        input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        input_ids = torch.tensor([[1, 2, 3, 4, 5]], device=self.device)
         attention_mask = torch.ones_like(input_ids)
 
         with torch.no_grad():
@@ -59,7 +60,7 @@ def test_mixture_coefficient(self):
         self.assertTrue(torch.allclose(wrapper_output.logits, expected_logits, atol=1e-5))
 
     def test_prepare_inputs_for_generation(self):
-        input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        input_ids = torch.tensor([[1, 2, 3, 4, 5]], device=self.device)
         attention_mask = torch.ones_like(input_ids)
 
         inputs = self.wrapper.prepare_inputs_for_generation(input_ids, attention_mask=attention_mask, use_cache=True)

tests/test_modeling_value_head.py

@@ -55,6 +55,10 @@ class VHeadModelTester(TrlTestCase):
         trl_model_class = None
         transformers_model_class = None
 
+        def setUp(self):
+            super().setUp()
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
         def test_value_head(self):
             r"""
             Test if the v-head is added to the model successfully
@@ -207,8 +211,8 @@ def test_inference(self):
         EXPECTED_OUTPUT_SIZE = 3
 
         for model_name in self.all_model_names:
-            model = self.trl_model_class.from_pretrained(model_name)
-            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
+            model = self.trl_model_class.from_pretrained(model_name).to(self.device)
+            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
             outputs = model(input_ids)
 
             # Check if the outputs are of the right size - here
@@ -250,8 +254,8 @@ def test_generate(self, model_name):
         Test if `generate` works for every model
         """
         generation_config = GenerationConfig(max_new_tokens=9)
-        model = self.trl_model_class.from_pretrained(model_name)
-        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
+        model = self.trl_model_class.from_pretrained(model_name).to(self.device)
+        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
 
         # Just check if the generation works
         _ = model.generate(input_ids, generation_config=generation_config)
@@ -263,7 +267,7 @@ def test_transformers_bf16_kwargs(self):
         run a dummy forward pass without any issue.
         """
         for model_name in self.all_model_names:
-            trl_model = self.trl_model_class.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+            trl_model = self.trl_model_class.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(self.device)
 
             lm_head_namings = ["lm_head", "embed_out", "output_layer"]
 
@@ -276,7 +280,7 @@ def test_transformers_bf16_kwargs(self):
                 if hasattr(trl_model.pretrained_model, lm_head_naming):
                     self.assertEqual(getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype, torch.bfloat16)
 
-            dummy_input = torch.LongTensor([[0, 1, 0, 1]])
+            dummy_input = torch.LongTensor([[0, 1, 0, 1]]).to(self.device)
 
             # check dummy forward pass works in half precision
             _ = trl_model(dummy_input)
@@ -323,9 +327,9 @@ def test_inference(self):
         EXPECTED_OUTPUT_SIZE = 3
 
         for model_name in self.all_model_names:
-            model = self.trl_model_class.from_pretrained(model_name)
-            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
-            decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
+            model = self.trl_model_class.from_pretrained(model_name).to(self.device)
+            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
+            decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
             outputs = model(input_ids, decoder_input_ids=decoder_input_ids)
 
             # Check if the outputs are of the right size - here
@@ -367,9 +371,9 @@ def test_generate(self, model_name):
         Test if `generate` works for every model
         """
         generation_config = GenerationConfig(max_new_tokens=9)
-        model = self.trl_model_class.from_pretrained(model_name)
-        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
-        decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
+        model = self.trl_model_class.from_pretrained(model_name).to(self.device)
+        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
+        decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
 
         # Just check if the generation works
         _ = model.generate(input_ids, decoder_input_ids=decoder_input_ids, generation_config=generation_config)
@@ -400,7 +404,7 @@ def test_transformers_bf16_kwargs(self):
         run a dummy forward pass without any issue.
         """
         for model_name in self.all_model_names:
-            trl_model = self.trl_model_class.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+            trl_model = self.trl_model_class.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(self.device)
 
             lm_head_namings = self.trl_model_class.lm_head_namings
 
@@ -412,7 +416,7 @@ def test_transformers_bf16_kwargs(self):
                 if hasattr(trl_model.pretrained_model, lm_head_naming):
                     self.assertTrue(getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype == torch.bfloat16)
 
-            dummy_input = torch.LongTensor([[0, 1, 0, 1]])
+            dummy_input = torch.LongTensor([[0, 1, 0, 1]]).to(self.device)
 
             # check dummy forward pass works in half precision
             _ = trl_model(input_ids=dummy_input, decoder_input_ids=dummy_input)

tests/test_utils.py

@@ -329,7 +329,8 @@ def setUp(self):
         super().setUp()
         # Initialize the tokenizer
         self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
-        self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = AutoModelForCausalLM.from_pretrained(self.model_id).to(self.device)
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
 
         self.generation_config = GenerationConfig(
@@ -350,7 +351,7 @@ def test_mini_batch_generation(self):
             self.tokenizer.apply_chat_template(example[:-1], add_generation_prompt=True, tokenize=False)
             for example in self.examples
         ]
-        queries = self.tokenizer(batch, padding=True, return_tensors="pt")["input_ids"]
+        queries = self.tokenizer(batch, padding=True, return_tensors="pt")["input_ids"].to(self.device)
         bs, context_length = queries.shape
 
         query_responses, logits = batch_generation(
@@ -369,7 +370,7 @@ def test_single_batch_generation(self):
             self.tokenizer.apply_chat_template(example[:-1], add_generation_prompt=True, tokenize=False)
             for example in self.examples
         ]
-        queries = self.tokenizer(batch, padding=True, return_tensors="pt")["input_ids"]
+        queries = self.tokenizer(batch, padding=True, return_tensors="pt")["input_ids"].to(self.device)
         bs, context_length = queries.shape
 
         query_responses, logits = batch_generation(

trl/trainer/dpo_trainer.py

@@ -1333,10 +1333,11 @@ def _compute_loss_liger(
                 model_kwargs["attention_mask"] = attention_mask
 
             # Get the base model outputs (before LM head)
-            if hasattr(unwrapped_model, "get_decoder"):
+            if hasattr(unwrapped_model, "get_decoder") and unwrapped_model.get_decoder() is not None:
                 base_model = unwrapped_model.get_decoder()
             else:
-                base_model = getattr(unwrapped_model, self.args.base_model_attribute_name, unwrapped_model)
+                base_attr = getattr(unwrapped_model, "base_model_prefix", self.args.base_model_attribute_name)
+                base_model = getattr(unwrapped_model, base_attr, unwrapped_model)
 
             outputs = base_model(
                 input_ids,
@@ -1349,12 +1350,11 @@ def _compute_loss_liger(
             ref_hidden_states = None
             if not self.reference_free and self.ref_model is not None:
                 unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model)
-                if hasattr(unwrapped_ref_model, "get_decoder"):
+                if hasattr(unwrapped_ref_model, "get_decoder") and unwrapped_ref_model.get_decoder() is not None:
                     ref_base_model = unwrapped_ref_model.get_decoder()
                 else:
-                    ref_base_model = getattr(
-                        unwrapped_ref_model, self.args.base_model_attribute_name, unwrapped_ref_model
-                    )
+                    ref_attr = getattr(unwrapped_ref_model, "base_model_prefix", self.args.base_model_attribute_name)
+                    ref_base_model = getattr(unwrapped_ref_model, ref_attr, unwrapped_ref_model)
 
                 ref_outputs = ref_base_model(
                     input_ids,
@@ -1363,10 +1363,11 @@ def _compute_loss_liger(
                 )
                 ref_hidden_states = ref_outputs.last_hidden_state[:, :-1]
             elif not self.reference_free:
-                if hasattr(unwrapped_model, "get_decoder"):
+                if hasattr(unwrapped_model, "get_decoder") and unwrapped_model.get_decoder() is not None:
                     ref_base_model = unwrapped_model.get_decoder()
                 else:
-                    ref_base_model = getattr(unwrapped_model, self.args.base_model_attribute_name, unwrapped_model)
+                    ref_attr = getattr(unwrapped_model, "base_model_prefix", self.args.base_model_attribute_name)
+                    ref_base_model = getattr(unwrapped_model, ref_attr, unwrapped_model)
                 with self.null_ref_context():
                     ref_outputs = ref_base_model(
                         input_ids,

trl/trainer/gkd_trainer.py

@@ -243,7 +243,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         if self.use_liger_gkd_loss:
             # Forward only through the base models (avoid lm_head to save memory)
             unwrapped_student = self.accelerator.unwrap_model(model)
-            if hasattr(unwrapped_student, "get_decoder"):
+            if hasattr(unwrapped_student, "get_decoder") and unwrapped_student.get_decoder() is not None:
                 base_student = unwrapped_student.get_decoder()
             else:
                 base_student = getattr(
@@ -259,7 +259,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
 
             self.teacher_model.eval()
             unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model)
-            if hasattr(unwrapped_teacher, "get_decoder"):
+            if hasattr(unwrapped_teacher, "get_decoder") and unwrapped_teacher.get_decoder() is not None:
                 base_teacher = unwrapped_teacher.get_decoder()
             else:
                 base_teacher = getattr(

trl/trainer/kto_trainer.py

@@ -1271,10 +1271,11 @@ def _compute_loss_liger(self, model, batch):
             )
         else:
             # skip the lm head and get the last hidden state
-            if hasattr(model, "get_decoder"):
+            if hasattr(model, "get_decoder") and model.get_decoder() is not None:
                 base_model = model.get_decoder()
             else:
-                base_model = getattr(model, self.args.base_model_attribute_name)
+                base_attr = getattr(model, "base_model_prefix", self.args.base_model_attribute_name)
+                base_model = getattr(model, base_attr, model)
             outputs = base_model(
                 batch["completion_input_ids"],
                 attention_mask=batch["completion_attention_mask"],
@@ -1283,10 +1284,11 @@ def _compute_loss_liger(self, model, batch):
             )
 
             # reference model
-            if hasattr(self.ref_model, "get_decoder"):
+            if hasattr(self.ref_model, "get_decoder") and self.ref_model.get_decoder() is not None:
                 ref_base_model = self.ref_model.get_decoder()
             else:
-                ref_base_model = getattr(self.ref_model, self.args.base_model_attribute_name)
+                ref_attr = getattr(self.ref_model, "base_model_prefix", self.args.base_model_attribute_name)
+                ref_base_model = getattr(self.ref_model, ref_attr, self.ref_model)
             ref_outputs = ref_base_model(
                 batch["completion_input_ids"],
                 attention_mask=batch["completion_attention_mask"],