Merge pull request #10 from philnach/fix/pytorch-inf-bfloat16-memory

Switch pytorch_inf from float32 to auto to fix OOM on 16GB machines

Author: jewilder

scenarios/macos/mac_pytorch_inf/mac_pytorch_inf.py

@@ -14,7 +14,7 @@
 class MacPytorchInf(core.app_scenario.Scenario):
 
     module = __module__.split('.')[-1]
-    prep_version = "5"
+    prep_version = "6"
     resources = module + "_resources"
 
 

scenarios/macos/mac_pytorch_inf/mac_pytorch_inf_resources/inference.py

@@ -185,7 +185,7 @@ def setup_model(model_name, device):
     print("Downloading model...")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        dtype=torch.float16 if device == 'cuda' else torch.float32,
+        torch_dtype="auto",
         device_map="auto" if device == 'cuda' else None
     )
     # model.resize_token_embeddings(len(tokenizer))
@@ -371,7 +371,7 @@ def main():
     print("Loading model...")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        dtype=torch.float16 if device == 'cuda' else torch.float32,
+        torch_dtype="auto",
         device_map="auto" if device == 'cuda' else None
     )
     model.resize_token_embeddings(len(tokenizer))

scenarios/windows/pytorch_inf/pytorch_inf.py

@@ -14,7 +14,7 @@
 class PytorchInf(core.app_scenario.Scenario):
 
     module = __module__.split('.')[-1]
-    prep_version = "9"
+    prep_version = "10"
     # prep_scenarios = [(module, prep_version)]
     resources = module + "_resources"
 

scenarios/windows/pytorch_inf/pytorch_inf_resources/inference.py

@@ -188,7 +188,7 @@ def setup_model(model_name, device):
     print("Downloading model...")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        dtype=torch.float16 if device == 'cuda' else torch.float32,
+        torch_dtype="auto",
         device_map="auto" if device == 'cuda' else None
     )
     # model.resize_token_embeddings(len(tokenizer))
@@ -374,7 +374,7 @@ def main():
     print("Loading model...")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        dtype=torch.float16 if device == 'cuda' else torch.float32,
+        torch_dtype="auto",
         device_map="auto" if device == 'cuda' else None
     )
     model.resize_token_embeddings(len(tokenizer))