SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

enable Pipeline to get device from model (#30534)

Browse Source 
* check model.device

* fix

* style fix

* move model device

* remove print

* add comment

* fix

* add unit test

* optimize

* change test names and add more cases

* Update tests/pipelines/test_pipelines_common.py

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

---------

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>
This commit is contained in:
Fanli Lin
2024-05-13 22:00:39 +08:00
committed by GitHub
parent f4dc26d466
commit 69d9bca55a
2 changed files with 51 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
tests/pipelines/test_pipelines_common.py
View File

@@ -48,6 +48,7 @@ from transformers.testing_utils import (
require_tf,
require_torch,
require_torch_accelerator,
require_torch_multi_accelerator,
require_torch_or_tf,
slow,
torch_device,
@@ -519,6 +520,52 @@ class PipelineUtilsTest(unittest.TestCase):
actual_output = classifier("Test input.")
self.assertEqual(expected_output, actual_output)
@require_torch_accelerator
def test_pipeline_no_device(self):
# Test when no device is passed to pipeline
import torch
from transformers import AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")
# Case 1: Model is manually moved to device
model = AutoModelForCausalLM.from_pretrained(
"hf-internal-testing/tiny-random-bert", torch_dtype=torch.float16
).to(torch_device)
model_device = model.device
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
self.assertEqual(pipe.model.device, model_device)
# Case 2: Model is loaded by accelerate
model = AutoModelForCausalLM.from_pretrained(
"hf-internal-testing/tiny-random-bert", device_map=torch_device, torch_dtype=torch.float16
)
model_device = model.device
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
self.assertEqual(pipe.model.device, model_device)
# Case 3: device_map is passed to model and device is passed to pipeline
model = AutoModelForCausalLM.from_pretrained(
"hf-internal-testing/tiny-random-bert", device_map=torch_device, torch_dtype=torch.float16
)
with self.assertRaises(ValueError):
pipe = pipeline("text-generation", model=model, device="cpu", tokenizer=tokenizer)
@require_torch_multi_accelerator
def test_pipeline_device_not_equal_model_device(self):
# Test when device ids are different, pipeline should move the model to the passed device id
import torch
from transformers import AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")
model_device = f"{torch_device}:1"
model = AutoModelForCausalLM.from_pretrained(
"hf-internal-testing/tiny-random-bert", torch_dtype=torch.float16
).to(model_device)
target_device = f"{torch_device}:0"
self.assertNotEqual(model_device, target_device)
pipe = pipeline("text-generation", model=model, device=target_device, tokenizer=tokenizer)
self.assertEqual(pipe.model.device, torch.device(target_device))
@slow
@require_torch
def test_load_default_pipelines_pt(self):