From b8f1cde931392551f74a9abef5d2724c3cbc2208 Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Mon, 16 Oct 2023 22:47:20 +0200 Subject: [PATCH] Fix Mistral OOM again (#26847) fix Co-authored-by: ydshieh --- tests/models/mistral/test_modeling_mistral.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/models/mistral/test_modeling_mistral.py b/tests/models/mistral/test_modeling_mistral.py index 311ed55892..a795ebceef 100644 --- a/tests/models/mistral/test_modeling_mistral.py +++ b/tests/models/mistral/test_modeling_mistral.py @@ -437,7 +437,8 @@ class MistralIntegrationTest(unittest.TestCase): input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338] model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto") input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device) - out = model(input_ids).logits.cpu() + with torch.no_grad(): + out = model(input_ids).logits.cpu() # Expected mean on dim = -1 EXPECTED_MEAN = torch.tensor([[-2.5548, -2.5737, -3.0600, -2.5906, -2.8478, -2.8118, -2.9325, -2.7694]]) torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2) @@ -457,8 +458,8 @@ class MistralIntegrationTest(unittest.TestCase): EXPECTED_TEXT_COMPLETION = """My favourite condiment is 100% ketchup. I love it on everything. I’m not a big""" prompt = "My favourite condiment is " tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", use_fast=False) - input_ids = tokenizer.encode(prompt, return_tensors="pt").to(torch_device) model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", device_map="auto") + input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.model.embed_tokens.weight.device) # greedy generation outputs generated_ids = model.generate(input_ids, max_new_tokens=20, temperature=0)