fix model path for int8 (#54)

Author: mayank31398

inference_server/models/ds_inference.py

@@ -6,13 +6,12 @@
 from functools import partial
 
 import torch
-import torch.distributed as dist
 
 import deepspeed
 from huggingface_hub import try_to_load_from_cache
 from transformers import AutoConfig
 
-from ..utils import get_world_size, print_rank_n, run_rank_n
+from ..utils import get_world_size, run_rank_n
 from .model import Model, get_hf_model_class
 
 
@@ -90,13 +89,17 @@ def __exit__(self, type, value, traceback):
 
 
 def get_model_path(model_name: str):
-    config_file = "config.json"
-
-    # will fall back to HUGGINGFACE_HUB_CACHE
-    config_path = try_to_load_from_cache(model_name, config_file, cache_dir=os.getenv("TRANSFORMERS_CACHE"))
-
-    if config_path is not None:
-        return os.path.dirname(config_path)
-    # treat the model name as an explicit model path
-    elif os.path.isfile(os.path.join(model_name, config_file)):
+    try:
+        config_file = "config.json"
+
+        # will fall back to HUGGINGFACE_HUB_CACHE
+        config_path = try_to_load_from_cache(model_name, config_file, cache_dir=os.getenv("TRANSFORMERS_CACHE"))
+
+        if config_path is None:
+            # treat the model name as an explicit model path
+            return model_name
+        elif os.path.isfile(os.path.join(model_name, config_file)):
+            return os.path.dirname(config_path)
+    except:
+        # treat the model name as an explicit model path
         return model_name

inference_server/models/ds_zero.py

@@ -1,8 +1,6 @@
-import os
 from argparse import Namespace
 
 import torch
-import torch.distributed as dist
 
 import deepspeed
 from transformers import AutoConfig

inference_server/models/hf_accelerate.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from ..utils import get_world_size, print_rank_n
+from ..utils import get_world_size
 from .model import Model, get_hf_model_class