From b5be72cb1f782e908fac33ba3bf180d88a45a76d Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Wed, 8 Feb 2023 15:14:56 +0000 Subject: [PATCH 1/2] add correct repo name --- tests/trainer/test_ppo_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/trainer/test_ppo_trainer.py b/tests/trainer/test_ppo_trainer.py index 3a59423c67..0e74f3b703 100644 --- a/tests/trainer/test_ppo_trainer.py +++ b/tests/trainer/test_ppo_trainer.py @@ -98,7 +98,7 @@ def tearDownClass(cls): def setUp(self): # model_id - model_id = "gpt2" + model_id = "trl-internal-testing/dummy-GPT2-correct-vocab" # get models and tokenizer self.gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(model_id) From 76c538305de11a0cbbf022b125dd49738b41016a Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Wed, 8 Feb 2023 15:29:11 +0000 Subject: [PATCH 2/2] fix tests --- tests/trainer/test_ppo_trainer.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/trainer/test_ppo_trainer.py b/tests/trainer/test_ppo_trainer.py index 0e74f3b703..8cae3d9db7 100644 --- a/tests/trainer/test_ppo_trainer.py +++ b/tests/trainer/test_ppo_trainer.py @@ -98,12 +98,12 @@ def tearDownClass(cls): def setUp(self): # model_id - model_id = "trl-internal-testing/dummy-GPT2-correct-vocab" + self.model_id = "trl-internal-testing/dummy-GPT2-correct-vocab" # get models and tokenizer - self.gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(model_id) - self.gpt2_model_ref = AutoModelForCausalLMWithValueHead.from_pretrained(model_id) - self.gpt2_tokenizer = GPT2Tokenizer.from_pretrained(model_id) + self.gpt2_model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id) + self.gpt2_model_ref = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id) + self.gpt2_tokenizer = GPT2Tokenizer.from_pretrained(self.model_id) # initialize trainer self.ppo_config = PPOConfig(batch_size=2, forward_batch_size=1, log_with=None) @@ -269,7 +269,7 @@ def test_ppo_step_with_no_ref(self): self.assertTrue(param.grad is None, f"Parameter {name} has a gradient") # initialize a new gpt2 model: - model = AutoModelForCausalLMWithValueHead.from_pretrained("gpt2") + model = AutoModelForCausalLMWithValueHead.from_pretrained(self.model_id) for name, param in ppo_trainer.ref_model.named_parameters(): if "v_head" not in name: name = name.replace("pretrained_model.", "") @@ -292,7 +292,7 @@ def test_ppo_step_with_no_ref_custom_layers(self): # initialize dataset dummy_dataset = self._init_dummy_dataset() - num_shared_layers = 6 + num_shared_layers = 1 ppo_trainer = PPOTrainer( config=self.ppo_config,