Bug fix for loading models with TP=1 when mii grpc server is not needed

kamalkraj · Apr 11, 2022 · e18abbc · e18abbc
1 parent 772087e
commit e18abbc
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 2 deletions.
diff --git a/examples/azure-local/gpt2-azure-local-example.py b/examples/azure-local/gpt2-azure-local-example.py
@@ -9,7 +9,7 @@
                resource_group=os.environ["AZ_RESOURCE_GROUP"])
 
 mii_configs = mii.constants.MII_CONFIGS_DEFAULT
-mii_configs[mii.constants.TENSOR_PARALLEL_KEY] = 2
+mii_configs[mii.constants.TENSOR_PARALLEL_KEY] = 1
 
 mii.deploy(task_name="text-generation",
            model_name="gpt2",

diff --git a/mii/__init__.py b/mii/__init__.py
@@ -7,3 +7,4 @@
 from .utils import setup_task, get_task, get_task_name, check_if_task_and_model_is_supported
 from .grpc_related.proto import modelresponse_pb2_grpc
 from .grpc_related.proto import modelresponse_pb2
+from .models.load_models import load_models
diff --git a/mii/server_client.py b/mii/server_client.py
@@ -112,7 +112,7 @@ def _is_server_process_alive(self):
     def _initialize_service(self, model_name, model_path, ds_optimize):
         process = None
         if not self.use_grpc_server:
-            self.model = mii.load_model(model_name, model_path)
+            self.model = mii.load_models(mii.get_task_name(self.task), model_name, model_path, ds_optimize)
         else:
             if self._is_socket_open(self.port_number):
                 raise RuntimeError(

diff --git a/mii/utils.py b/mii/utils.py
@@ -143,6 +143,8 @@ def setup_task():
     return get_model_path(), not is_aml(), is_aml()
 
 
+
+
 log_levels = {
     "debug": logging.DEBUG,
     "info": logging.INFO,