Skip to content

Commit ddf1cb2

Browse files
albertvillanovanenuadrian
albertvillanova
authored and
nenuadrian
committed
Set model dtype to float32 in experimental tests of trainers (huggingface#4925)
1 parent 254cd81 commit ddf1cb2

11 files changed

Lines changed: 34 additions & 30 deletions

‎tests/experimental/test_bco_trainer.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ class TestBCOTrainer(TrlTestCase):
4747
@require_sklearn
4848
def test_train(self, config_name):
4949
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
50-
model = AutoModelForCausalLM.from_pretrained(model_id)
50+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
5151
ref_model = AutoModelForCausalLM.from_pretrained(model_id)
5252
tokenizer = AutoTokenizer.from_pretrained(model_id)
5353

@@ -83,7 +83,7 @@ def test_train(self, config_name):
8383
@require_sklearn
8484
def test_train_with_precompute(self):
8585
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
86-
model = AutoModelForCausalLM.from_pretrained(model_id)
86+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
8787
ref_model = AutoModelForCausalLM.from_pretrained(model_id)
8888
tokenizer = AutoTokenizer.from_pretrained(model_id)
8989

@@ -120,7 +120,7 @@ def test_train_with_precompute(self):
120120
@require_sklearn
121121
def test_train_eval(self):
122122
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
123-
model = AutoModelForCausalLM.from_pretrained(model_id)
123+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
124124
ref_model = AutoModelForCausalLM.from_pretrained(model_id)
125125
tokenizer = AutoTokenizer.from_pretrained(model_id)
126126

@@ -148,7 +148,7 @@ def test_train_eval(self):
148148
@require_sklearn
149149
def test_init_with_ref_model_is_model(self):
150150
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
151-
model = AutoModelForCausalLM.from_pretrained(model_id)
151+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
152152
tokenizer = AutoTokenizer.from_pretrained(model_id)
153153

154154
dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train")
@@ -171,7 +171,7 @@ def test_init_with_ref_model_is_model(self):
171171
@require_sklearn
172172
def test_tokenize_and_process_tokens(self):
173173
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
174-
model = AutoModelForCausalLM.from_pretrained(model_id)
174+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
175175
ref_model = AutoModelForCausalLM.from_pretrained(model_id)
176176
tokenizer = AutoTokenizer.from_pretrained(model_id)
177177

@@ -226,7 +226,7 @@ def test_tokenize_and_process_tokens(self):
226226
@require_sklearn
227227
def test_train_without_providing_ref_model(self):
228228
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
229-
model = AutoModelForCausalLM.from_pretrained(model_id)
229+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
230230
tokenizer = AutoTokenizer.from_pretrained(model_id)
231231

232232
dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train")
@@ -260,7 +260,7 @@ def test_train_without_providing_ref_model(self):
260260
@require_sklearn
261261
def test_train_udm(self):
262262
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
263-
model = AutoModelForCausalLM.from_pretrained(model_id)
263+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
264264
tokenizer = AutoTokenizer.from_pretrained(model_id)
265265

266266
# Get embedding model
@@ -310,7 +310,7 @@ def embed_prompt(input_ids, attention_mask, model):
310310
@require_peft
311311
def test_train_without_providing_ref_model_with_lora(self):
312312
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
313-
model = AutoModelForCausalLM.from_pretrained(model_id)
313+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
314314
lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM")
315315
tokenizer = AutoTokenizer.from_pretrained(model_id)
316316

@@ -348,7 +348,7 @@ def test_train_without_providing_ref_model_with_lora(self):
348348
@require_no_wandb
349349
def test_generate_during_eval_no_wandb(self):
350350
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
351-
model = AutoModelForCausalLM.from_pretrained(model_id)
351+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
352352
tokenizer = AutoTokenizer.from_pretrained(model_id)
353353

354354
dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference")
@@ -379,7 +379,7 @@ def test_generate_during_eval_no_wandb(self):
379379
@require_peft
380380
def test_lora_train_and_save(self):
381381
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
382-
model = AutoModelForCausalLM.from_pretrained(model_id)
382+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
383383
lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM")
384384
tokenizer = AutoTokenizer.from_pretrained(model_id)
385385

@@ -411,7 +411,7 @@ def test_lora_train_and_save(self):
411411
@require_sklearn
412412
def test_compute_metrics(self):
413413
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
414-
model = AutoModelForCausalLM.from_pretrained(model_id)
414+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
415415
ref_model = AutoModelForCausalLM.from_pretrained(model_id)
416416
tokenizer = AutoTokenizer.from_pretrained(model_id)
417417

‎tests/experimental/test_cpo_trainer.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,13 @@
2525
class TestCPOTrainer(TrlTestCase):
2626
def setup_method(self):
2727
self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
28-
self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
28+
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32")
2929
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
3030
self.tokenizer.pad_token = self.tokenizer.eos_token
3131

3232
# get t5 as seq2seq example:
3333
model_id = "trl-internal-testing/tiny-T5ForConditionalGeneration"
34-
self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
34+
self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id, dtype="float32")
3535
self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)
3636

3737
@pytest.mark.parametrize(

‎tests/experimental/test_gkd_trainer.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def setup_class(cls):
3232
cls.device = "cuda" if torch.cuda.is_available() else "cpu"
3333
cls.tokenizer = AutoTokenizer.from_pretrained(model_id)
3434
cls.tokenizer.pad_token = cls.tokenizer.eos_token
35-
cls.model = AutoModelForCausalLM.from_pretrained(model_id).to(cls.device)
35+
cls.model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32").to(cls.device)
3636
cls.generation_config = GenerationConfig(
3737
max_new_tokens=20,
3838
num_return_sequences=1,
@@ -201,7 +201,7 @@ def test_zero_loss_for_identical_inputs(self):
201201
class TestGKDTrainer(TrlTestCase):
202202
def setup_method(self):
203203
self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
204-
self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
204+
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32")
205205
self.teacher_model = AutoModelForCausalLM.from_pretrained(self.model_id)
206206
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
207207
self.tokenizer.pad_token = self.tokenizer.eos_token

‎tests/experimental/test_kto_trainer.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -314,7 +314,7 @@ def test_kto_trainer_with_liger(self):
314314
assert not torch.equal(param, new_param)
315315

316316
def test_compute_metrics(self):
317-
model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
317+
model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", dtype="float32")
318318
ref_model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
319319
tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
320320
tokenizer.pad_token = tokenizer.eos_token

‎tests/experimental/test_merge_model_callback.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@
2828
@require_mergekit
2929
class TestMergeModelCallback(TrlTestCase):
3030
def setup_method(self):
31-
self.model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
31+
self.model = AutoModelForCausalLM.from_pretrained(
32+
"trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", dtype="float32"
33+
)
3234
self.tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
3335
self.dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")
3436

‎tests/experimental/test_nash_md_trainer.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ class TestGeometricMixtureWrapper(TrlTestCase):
3434
def setup_method(self):
3535
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
3636
self.device = "cuda" if torch.cuda.is_available() else "cpu"
37-
self.model = AutoModelForCausalLM.from_pretrained(model_id).to(self.device)
37+
self.model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32").to(self.device)
3838
self.ref_model = create_reference_model(self.model).to(self.device)
3939
self.generation_config = GenerationConfig.from_pretrained(model_id)
4040
self.mixture_coef = 0.5
@@ -81,7 +81,7 @@ def test_prepare_inputs_for_generation(self):
8181
class TestNashMDTrainer(TrlTestCase):
8282
def setup_method(self):
8383
self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
84-
self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
84+
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32")
8585
self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)
8686
self.reward_model = AutoModelForSequenceClassification.from_pretrained(self.model_id, num_labels=1)
8787
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)

‎tests/experimental/test_online_dpo_trainer.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
class TestOnlineDPOTrainer(TrlTestCase):
4545
def setup_method(self):
4646
self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
47-
self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
47+
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32")
4848
self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)
4949
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
5050
self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -251,7 +251,7 @@ def cleanup_vllm_communicator(trainer):
251251
pass # Continue if cleanup fails
252252

253253
model_id = "trl-internal-testing/small-Qwen2ForCausalLM-2.5" # We need a bigger model
254-
model = AutoModelForCausalLM.from_pretrained(model_id)
254+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
255255
tokenizer = AutoTokenizer.from_pretrained(model_id)
256256
tokenizer.pad_token = tokenizer.eos_token
257257

@@ -284,7 +284,7 @@ def cleanup_vllm_communicator(trainer):
284284
def test_training_with_vllm_colocate(self):
285285
"""Test vLLM colocate mode with our refactored implementation"""
286286
model_id = "trl-internal-testing/small-Qwen2ForCausalLM-2.5" # We need a bigger model
287-
model = AutoModelForCausalLM.from_pretrained(model_id)
287+
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32")
288288
tokenizer = AutoTokenizer.from_pretrained(model_id)
289289
tokenizer.pad_token = tokenizer.eos_token
290290

@@ -490,7 +490,7 @@ def test_online_dpo_vlm_trainer(self, model_id):
490490
dataset = Dataset.from_dict(dataset_dict)
491491
dataset = dataset.cast_column("images", features.Sequence(features.Image()))
492492

493-
model = AutoModelForImageTextToText.from_pretrained(model_id)
493+
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype="float32")
494494
reward_model = AutoModelForSequenceClassification.from_pretrained(
495495
"trl-internal-testing/tiny-LlamaForCausalLM-3.2", num_labels=1
496496
)

‎tests/experimental/test_orpo_trainer.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,13 @@
2525
class TestORPOTrainer(TrlTestCase):
2626
def setup_method(self):
2727
self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
28-
self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
28+
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32")
2929
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
3030
self.tokenizer.pad_token = self.tokenizer.eos_token
3131

3232
# get t5 as seq2seq example:
3333
model_id = "trl-internal-testing/tiny-T5ForConditionalGeneration"
34-
self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
34+
self.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id, dtype="float32")
3535
self.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)
3636

3737
@pytest.mark.parametrize(
@@ -143,7 +143,7 @@ def test_orpo_trainer_with_lora(self, config_name):
143143
assert not torch.equal(param, new_param)
144144

145145
def test_compute_metrics(self):
146-
model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
146+
model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", dtype="float32")
147147
tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
148148
tokenizer.pad_token = tokenizer.eos_token
149149

‎tests/experimental/test_ppo_trainer.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ def setup_method(self):
7676
# Initialize the tokenizer
7777
self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
7878
self.device = "cuda" if torch.cuda.is_available() else "cpu"
79-
self.model = AutoModelForCausalLM.from_pretrained(self.model_id).to(self.device)
79+
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32").to(self.device)
8080
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
8181

8282
self.generation_config = GenerationConfig(
@@ -688,7 +688,7 @@ class TestPPOTrainer(TrlTestCase):
688688
def setup_method(self):
689689
# Set up the models and tokenizer using the test model
690690
self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
691-
self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
691+
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype="float32")
692692
self.ref_model = AutoModelForCausalLM.from_pretrained(self.model_id)
693693
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, padding_side="left")
694694
self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})

‎tests/experimental/test_winrate_callback.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,9 @@ def __init__(self, model, ref_model, args, train_dataset, eval_dataset, processi
5454

5555
class TestWinRateCallback(TrlTestCase):
5656
def setup_method(self):
57-
self.model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
57+
self.model = AutoModelForCausalLM.from_pretrained(
58+
"trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", dtype="float32"
59+
)
5860
self.ref_model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
5961
self.tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5")
6062
self.tokenizer.pad_token = self.tokenizer.eos_token

0 commit comments

Comments
 (0)