H100 adaptability

TUDB-Labs · Nov 28, 2024 · 10bad92 · 10bad92
1 parent 26d2e1e
commit 10bad92
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/mlora/model/llm/model_llama.py b/mlora/model/llm/model_llama.py
@@ -217,6 +217,10 @@ def create_device_map() -> str | Dict[str, str]:
             "torch_dtype": torch.float32,
         }
 
+        # If it is an H100 device, set precision to nf4
+        if torch.cuda.is_available() and "H100" in torch.cuda.get_device_name(torch.cuda.current_device()):
+            precision = "nf4"
+
         logging.info(f"Loading model with precision - {precision}")
 
         if precision in load_type_dict: