riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-04 05:10:06 +06:00
Code Issues Actions 20 Packages Projects Releases Wiki Activity

enable test_assisted_decoding_in_different_gpu test on XPU (#37120)

Browse Source 
Signed-off-by: YAO Matrix <matrix.yao@intel.com>
Co-authored-by: Yih-Dar <2521628+ydshieh@users.noreply.github.com>
This commit is contained in:
Yao Matrix 2025-04-01 17:22:59 +08:00 committed by GitHub
parent 737cbd2109
commit 8f6b27eb5c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tests/generation/test_utils.py
View File

@ -3748,11 +3748,13 @@ class GenerationIntegrationTests(unittest.TestCase):
self.assertTrue(y_prob <= 1.0 and n_prob <= 1.0)
@slow
@require_torch_multi_gpu
@require_torch_multi_accelerator
def test_assisted_decoding_in_different_gpu(self):
model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to("cuda:0")
device_0 = f"{torch_device}:0" if torch_device != "cpu" else "cpu"
device_1 = f"{torch_device}:1" if torch_device != "cpu" else "cpu"
model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to(device_0)
assistant = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM").to(
"cuda:1"
device_1
)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM")
model.config.pad_token_id = tokenizer.eos_token_id