enable support for device map auto

Author: tarsur909

README.md

@@ -207,7 +207,7 @@ export HF_ACCESS_TOKEN="your_huggingface_api_key"
 
 - `task_id` (int, optional): Problem task id for selecting a problem from a Dataset.
 
-- `use_auto_device_mapping` (bool, optional): Use auto device mapping. Defaults to False.
+- `device_map` (str, optional): Device map for the model. Defaults to None.
 - `kwargs`(void, optional): Currently supported `kwargs` are `max_length`, `max_new_tokens`, `min_length`, `min_new_tokens`, `early_stopping`, `do_sample`, `num_beams`, `use_cache`, `temperature`, `top_k`, `top_p`, `num_return_sequences`, `pad_token_id`, and `eos_token_id`. Refer to the [HuggingFace Text Generation Documentation](https://huggingface.co/docs/transformers/en/main_classes/text_generation) for more information.
 
 
@@ -238,7 +238,7 @@ python3 syncode/infer.py
     --new_mask_store [True, False]
     --parser ["lr", "lalr"]
     --task_id [task_id]
-    --use_auto_device_mapping [True, False]
+    --device_map [device_map]
 ```
 </details>
 

syncode/common.py

@@ -11,17 +11,17 @@
 HF_ACCESS_TOKEN = os.environ['HF_ACCESS_TOKEN'] if 'HF_ACCESS_TOKEN' in os.environ else None
 
 
-def load_model(model_name, device, quantize, use_auto_device_mapping = False):
+def load_model(model_name, device, quantize, device_map = None):
         if model_name == 'test':
             model = AutoModelForCausalLM.from_pretrained('bigcode/tiny_starcoder_py').to(device)
         elif model_name == 'test-instruct':
             model = AutoModelForCausalLM.from_pretrained("rahuldshetty/tiny-starcoder-instruct")
         else:
-            if use_auto_device_mapping:
+            if device_map is not None:
                 if (quantize):
-                    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, cache_dir=HF_CACHE, token=HF_ACCESS_TOKEN, trust_remote_code=True, device_map = 'auto').eval()
+                    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, cache_dir=HF_CACHE, token=HF_ACCESS_TOKEN, trust_remote_code=True, device_map = device_map).eval()
                 else:
-                    model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=HF_CACHE, token=HF_ACCESS_TOKEN, trust_remote_code=True, device_map = 'auto').eval()
+                    model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=HF_CACHE, token=HF_ACCESS_TOKEN, trust_remote_code=True, device_map = device_map).eval()
             else:
                 if (quantize):
                     model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, cache_dir=HF_CACHE, token=HF_ACCESS_TOKEN, trust_remote_code=True).eval().to(device)

syncode/infer.py

@@ -50,7 +50,7 @@ def __init__(
         parser: Literal["lr", "lalr"] = "lalr",
         seed: Optional[int] = None,
         opp: bool = True,
-        use_auto_device_mapping: bool = False,
+        device_map: Optional[str] = None,
         **kwargs
     ):  
         # Check inputs
@@ -86,7 +86,7 @@ def __init__(
         self.grammar = Grammar(grammar) if self._is_grammar_mode() else None
 
         # Load model and tokenizer
-        model = common.load_model(self.model_name, device, quantize, use_auto_device_mapping)
+        model = common.load_model(self.model_name, device, quantize, device_map)
         tokenizer = common.load_tokenizer(self.model_name)
 
         # Initialize grammar decoder if needed
@@ -260,7 +260,7 @@ def main(
     parse_output_only: bool = True,
     prompt_type: str = 'original',
     format_tabs: bool = False,
-    use_auto_device_mapping: bool = False,
+    device_map: Optional[str] = None,
     **kwargs
 ):
     """Run Syncode with the specified configuration.
@@ -311,7 +311,7 @@ def main(
         seed=seed, 
         opp=opp,
         parse_output_only=parse_output_only,
-        use_auto_device_mapping=use_auto_device_mapping,
+        device_map=device_map,
         **kwargs
     )