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

Fix GPU OOM for mistral.py::Mask4DTestHard (#31212)

Browse Source 
* build

* build

* build

* build

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2024-06-03 19:25:15 +02:00
committed by GitHub
parent df5abae894
commit 8a1a23ae4d
1 changed files with 13 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/models/mistral/test_modeling_mistral.py
View File

@@ -734,15 +734,24 @@ class MistralIntegrationTest(unittest.TestCase):
@slow @slow
@require_torch_gpu @require_torch_gpu
class Mask4DTestHard(unittest.TestCase): class Mask4DTestHard(unittest.TestCase):
model_name = "mistralai/Mistral-7B-v0.1"
_model = None
def tearDown(self): def tearDown(self):
gc.collect() gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()
@property
def model(self):
if self.__class__._model is None:
self.__class__._model = MistralForCausalLM.from_pretrained(
self.model_name, torch_dtype=self.model_dtype
).to(torch_device)
return self.__class__._model
def setUp(self): def setUp(self):
model_name = "mistralai/Mistral-7B-v0.1" self.model_dtype = torch.float16
self.model_dtype = torch.float32 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False)
self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
self.model = MistralForCausalLM.from_pretrained(model_name, torch_dtype=self.model_dtype).to(torch_device)
def get_test_data(self): def get_test_data(self):
template = "my favorite {}" template = "my favorite {}"