Read HF device map from model config

danielz02 · Oct 12, 2023 · a070ee1 · a070ee1
1 parent e6c9452
commit a070ee1
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/src/helm/proxy/clients/huggingface_client.py b/src/helm/proxy/clients/huggingface_client.py
@@ -67,7 +67,7 @@ def __init__(self, model_config: HuggingFaceModelConfig):
         elif model_config.model_dtype == WeightType.BFLOAT16:
             model_kwargs["torch_dtype"] = torch.bfloat16
 
-        model_kwargs["device_map"] = "auto"
+        model_kwargs["device_map"] = model_config.device_map
 
         with htrack_block(f"Loading Hugging Face model for config {model_config}"):
             # WARNING this may fail if your GPU does not have enough memory

diff --git a/src/helm/proxy/clients/huggingface_model_registry.py b/src/helm/proxy/clients/huggingface_model_registry.py
@@ -39,8 +39,6 @@ class HuggingfaceModelQuantizationConfig:
 
     disable_exllama: bool = True
 
-    device_map: str = "auto"
-
     quant_file: Optional[str] = None
     """Path to .pt file generated by AWQ. This argument is required when using AutoAWQ"""