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

[bnb] fix bnb failing test (#22439)

Browse Source 
* fix bnb failing test

* fix

* fix

* fixup
This commit is contained in:
Younes Belkada
2023-03-29 15:13:00 +02:00
committed by GitHub
parent fab1de72f1
commit 33f4cb1093
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tests/mixed_int8/test_mixed_int8.py
View File

@@ -384,9 +384,8 @@ class MixedInt8TestMultiGpu(BaseMixedInt8Test):
Let's just try to load a model on 2 GPUs and see if it works. The model we test has ~2GB of total, 3GB should suffice Let's just try to load a model on 2 GPUs and see if it works. The model we test has ~2GB of total, 3GB should suffice
""" """
memory_mapping = {0: "1GB", 1: "2GB"}
model_parallel = AutoModelForCausalLM.from_pretrained( model_parallel = AutoModelForCausalLM.from_pretrained(
self.model_name, load_in_8bit=True, max_memory=memory_mapping, device_map="auto" self.model_name, load_in_8bit=True, device_map="balanced"
) )
# Check correct device map # Check correct device map