From 9d7afd2536ecd9816dd2ea9592a01e52fec17d17 Mon Sep 17 00:00:00 2001 From: Alex McKinney <44398246+vvvm23@users.noreply.github.com> Date: Fri, 18 Aug 2023 11:40:40 +0100 Subject: [PATCH] Replaces calls to `.cuda` with `.to(torch_device)` in tests (#25571) * Replaces calls to `.cuda` with `.to(torch_device)` in tests `torch.Tensor.cuda()` is a pre-0.4 solution to changing a tensor's device. It is recommended to prefer `.to(...)` for greater flexibility and error handling. Furthermore, this makes it more consistent with other tests (that tend to use `.to(torch_device)`) and ensures the correct device backend is used (if `torch_device` is neither `cpu` or `cuda`). * addressing review comments * more formatting changes in Bloom test * `make style` * Update tests/models/bloom/test_modeling_bloom.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * fixes style failures --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- tests/models/bloom/test_modeling_bloom.py | 25 ++++++++++--------- tests/models/jukebox/test_modeling_jukebox.py | 18 +++++++------ tests/models/opt/test_modeling_opt.py | 6 ++--- tests/models/xglm/test_modeling_xglm.py | 6 ++--- 4 files changed, 29 insertions(+), 26 deletions(-) diff --git a/tests/models/bloom/test_modeling_bloom.py b/tests/models/bloom/test_modeling_bloom.py index de7cb03e7e..d53362c997 100644 --- a/tests/models/bloom/test_modeling_bloom.py +++ b/tests/models/bloom/test_modeling_bloom.py @@ -423,7 +423,7 @@ class BloomModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi # >=1b1 + allow_fp16_reduced_precision_reduction = False + torch.bmm ==> PASS path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").cuda() + model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) model = model.eval() tokenizer = BloomTokenizerFast.from_pretrained(path_560m) @@ -435,7 +435,7 @@ class BloomModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi ) input_ids = tokenizer.encode(input_sentence, return_tensors="pt") - greedy_output = model.generate(input_ids.cuda(), max_length=50) + greedy_output = model.generate(input_ids.to(torch_device), max_length=50) self.assertEqual(tokenizer.decode(greedy_output[0], skip_special_tokens=True), EXPECTED_OUTPUT) @@ -443,16 +443,16 @@ class BloomModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi @require_torch_gpu def test_batch_generation(self): path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").cuda() + model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) model = model.eval() tokenizer = BloomTokenizerFast.from_pretrained(path_560m, padding_side="left") input_sentence = ["I enjoy walking with my cute dog", "I enjoy walking with my cute dog"] input_ids = tokenizer.batch_encode_plus(input_sentence, return_tensors="pt", padding=True) - greedy_output = model.generate( - input_ids["input_ids"].cuda(), attention_mask=input_ids["attention_mask"], max_length=50, do_sample=False - ) + input_ids = input_ids["input_ids"].to(torch_device) + attention_mask = input_ids["attention_mask"] + greedy_output = model.generate(input_ids, attention_mask=attention_mask, max_length=50, do_sample=False) self.assertEqual( tokenizer.decode(greedy_output[0], skip_special_tokens=True), @@ -463,7 +463,7 @@ class BloomModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi @require_torch_gpu def test_batch_generation_padd(self): path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").cuda() + model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) model = model.eval() tokenizer = BloomTokenizerFast.from_pretrained(path_560m, padding_side="left") @@ -473,10 +473,11 @@ class BloomModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi input_ids = tokenizer.batch_encode_plus(input_sentence, return_tensors="pt", padding=True) input_ids_without_pad = tokenizer.encode(input_sentence_without_pad, return_tensors="pt") - greedy_output = model.generate( - input_ids["input_ids"].cuda(), attention_mask=input_ids["attention_mask"], max_length=50, do_sample=False + input_ids, attention_mask = input_ids["input_ids"].to(torch_device), input_ids["attention_mask"] + greedy_output = model.generate(input_ids, attention_mask=attention_mask, max_length=50, do_sample=False) + greedy_output_without_pad = model.generate( + input_ids_without_pad.to(torch_device), max_length=50, do_sample=False ) - greedy_output_without_pad = model.generate(input_ids_without_pad.cuda(), max_length=50, do_sample=False) # test token values self.assertEqual(greedy_output[-1, 3:].tolist(), greedy_output_without_pad[0, :-3].tolist()) @@ -492,7 +493,7 @@ class BloomModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi def test_batch_generated_text(self): path_560m = "bigscience/bloom-560m" - model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").cuda() + model = BloomForCausalLM.from_pretrained(path_560m, use_cache=True, revision="gs555750").to(torch_device) model = model.eval() tokenizer = BloomTokenizerFast.from_pretrained(path_560m, padding_side="left") @@ -502,7 +503,7 @@ class BloomModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi ] inputs = tokenizer(input_sentences, return_tensors="pt", padding=True, truncation=True) generated_ids = model.generate( - inputs["input_ids"].cuda(), attention_mask=inputs["attention_mask"], max_length=20 + inputs["input_ids"].to(torch_device), attention_mask=inputs["attention_mask"], max_length=20 ) generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) diff --git a/tests/models/jukebox/test_modeling_jukebox.py b/tests/models/jukebox/test_modeling_jukebox.py index c0c78a25f8..8de0696c04 100644 --- a/tests/models/jukebox/test_modeling_jukebox.py +++ b/tests/models/jukebox/test_modeling_jukebox.py @@ -16,7 +16,7 @@ import unittest from unittest import skip from transformers import is_torch_available -from transformers.testing_utils import require_torch, slow +from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device from transformers.trainer_utils import set_seed @@ -363,35 +363,37 @@ class Jukebox5bModelTester(unittest.TestCase): self.assertIn(zs[2][0].detach().cpu().tolist(), [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]) @slow + @require_torch_gpu @skip("Not enough GPU memory on CI runners") def test_slow_sampling(self): model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval() - labels = [i.cuda() for i in self.prepare_inputs(self.model_id)] + labels = [i.to(torch_device) for i in self.prepare_inputs(self.model_id)] set_seed(0) - model.priors[0].cuda() - zs = [torch.zeros(1, 0, dtype=torch.long).cuda() for _ in range(3)] + model.priors[0].to(torch_device) + zs = [torch.zeros(1, 0, dtype=torch.long).to(torch_device) for _ in range(3)] zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False) torch.testing.assert_allclose(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2)) model.priors[0].cpu() set_seed(0) - model.priors[1].cuda() + model.priors[1].to(torch_device) zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False) torch.testing.assert_allclose(zs[1][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_1)) model.priors[1].cpu() set_seed(0) - model.priors[2].cuda() + model.priors[2].to(torch_device) zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False) torch.testing.assert_allclose(zs[2][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_0)) @slow + @require_torch_gpu def test_fp16_slow_sampling(self): prior_id = "ArthurZ/jukebox_prior_0" - model = JukeboxPrior.from_pretrained(prior_id, min_duration=0).eval().half().to("cuda") + model = JukeboxPrior.from_pretrained(prior_id, min_duration=0).eval().half().to(torch_device) - labels = self.prepare_inputs(prior_id)[0].cuda() + labels = self.prepare_inputs(prior_id)[0].to(torch_device) metadata = model.get_metadata(labels, 0, 7680, 0) set_seed(0) outputs = model.sample(1, metadata=metadata, sample_tokens=60) diff --git a/tests/models/opt/test_modeling_opt.py b/tests/models/opt/test_modeling_opt.py index 669dc3e6c0..232a29d4f5 100644 --- a/tests/models/opt/test_modeling_opt.py +++ b/tests/models/opt/test_modeling_opt.py @@ -522,13 +522,13 @@ class OPTGenerationTest(unittest.TestCase): model_name = "facebook/opt-1.3b" tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_fast=False, padding_side="left") - model = OPTForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, use_cache=True).cuda() + model = OPTForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, use_cache=True).to(torch_device) model = model.eval() batch = tokenizer(["Who are you?", "Joe Biden is the president of"], padding=True, return_tensors="pt") - input_ids = batch["input_ids"].cuda() - attention_mask = batch["attention_mask"].cuda() + input_ids = batch["input_ids"].to(torch_device) + attention_mask = batch["attention_mask"].to(torch_device) with torch.no_grad(): outputs = model(input_ids, attention_mask=attention_mask) diff --git a/tests/models/xglm/test_modeling_xglm.py b/tests/models/xglm/test_modeling_xglm.py index 9cf0c9a2bb..105ad5c44e 100644 --- a/tests/models/xglm/test_modeling_xglm.py +++ b/tests/models/xglm/test_modeling_xglm.py @@ -497,13 +497,13 @@ class XGLMModelLanguageGenerationTest(unittest.TestCase): model_name = "facebook/xglm-564M" tokenizer = XGLMTokenizer.from_pretrained(model_name, use_fast=False, padding_side="left") - model = XGLMForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, use_cache=True).cuda() + model = XGLMForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, use_cache=True).to(torch_device) model = model.eval() batch = tokenizer(["Who are you?", "Joe Biden is the president of"], padding=True, return_tensors="pt") - input_ids = batch["input_ids"].cuda() - attention_mask = batch["attention_mask"].cuda() + input_ids = batch["input_ids"].to(torch_device) + attention_mask = batch["attention_mask"].to(torch_device) with torch.no_grad(): outputs = model(input_ids, attention_mask=attention_mask)