Log device_map in HELM

danielz02 · Oct 28, 2023 · 3a2b0d3 · 3a2b0d3
1 parent b88b36d
commit 3a2b0d3
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/helm/proxy/clients/huggingface_client.py b/src/helm/proxy/clients/huggingface_client.py
@@ -87,10 +87,13 @@ def __init__(self, model_config: HuggingFaceModelConfig):
                     inject_fused_attention=model_config.quantization_config.inject_fused_attention,
                     **model_kwargs
                 )
+                hlog(self.model.hf_device_map)
             else:
                 self.model = AutoModelForCausalLM.from_pretrained(
                     model_name, trust_remote_code=True, **model_kwargs
                 )
+                hlog(self.model.hf_device_map)
+
         with htrack_block(f"Loading Hugging Face tokenizer model for config {model_config}"):
             # When the quantized model has uses a different tokenizer than its moddel name
             if model_config.tokenizer_name: