add eager-mode to cli (#2645)

benchmark/profile_generation.py

@@ -341,6 +341,8 @@ def parse_args():
     ArgumentHelper.backend(parser)
     # pytorch engine args
     pt_group = parser.add_argument_group('PyTorch engine arguments')
+    ArgumentHelper.eager_mode(pt_group)
+
     tp_act = ArgumentHelper.tp(pt_group)
     cache_count_act = ArgumentHelper.cache_max_entry_count(pt_group)
     cache_block_seq_len_act = ArgumentHelper.cache_block_seq_len(pt_group)
@@ -422,6 +424,7 @@ def main():
                     session_len=session_len,
                     tp=args.tp,
                     thread_safe=True,
+                    eager_mode=args.eager_mode,
                     enable_prefix_caching=args.enable_prefix_caching,
                 )
             gen_config = GenerationConfig(top_k=args.top_k,

benchmark/profile_pipeline_api.py

@@ -194,6 +194,8 @@ def parse_args():
 
     # pytorch engine args
     pt_group = parser.add_argument_group('PyTorch engine arguments')
+    ArgumentHelper.eager_mode(pt_group)
+
     tp_act = ArgumentHelper.tp(pt_group)
     session_len_act = ArgumentHelper.session_len(pt_group, default=4096)
     cache_count_act = ArgumentHelper.cache_max_entry_count(pt_group)
@@ -241,6 +243,7 @@ def main():
             max_batch_size=args.concurrency,
             tp=args.tp,
             thread_safe=False,
+            eager_mode=args.eager_mode,
             enable_prefix_caching=args.enable_prefix_caching,
         )
 

benchmark/profile_throughput.py

@@ -281,6 +281,8 @@ def parse_args():
 
     # pytorch engine args
     pt_group = parser.add_argument_group('PyTorch engine arguments')
+    ArgumentHelper.eager_mode(pt_group)
+
     tp_act = ArgumentHelper.tp(pt_group)
     session_len_act = ArgumentHelper.session_len(pt_group, default=4096)
     cache_count_act = ArgumentHelper.cache_max_entry_count(pt_group)
@@ -328,6 +330,7 @@ def main():
             max_batch_size=args.concurrency,
             tp=args.tp,
             thread_safe=True,
+            eager_mode=args.eager_mode,
             enable_prefix_caching=args.enable_prefix_caching,
             quant_policy=args.quant_policy,
         )

lmdeploy/cli/cli.py

@@ -122,6 +122,7 @@ def add_parser_chat():
         pt_group = parser.add_argument_group('PyTorch engine arguments')
         ArgumentHelper.adapters(pt_group)
         ArgumentHelper.device(pt_group)
+        ArgumentHelper.eager_mode(pt_group)
         # common engine args
         dtype_act = ArgumentHelper.dtype(pt_group)
         tp_act = ArgumentHelper.tp(pt_group)
@@ -265,6 +266,7 @@ def chat(args):
                 adapters=adapters,
                 enable_prefix_caching=args.enable_prefix_caching,
                 device_type=args.device,
+                eager_mode=args.eager_mode,
                 quant_policy=args.quant_policy)
             run_chat(args.model_path,
                      engine_config,
@@ -275,6 +277,7 @@ def chat(args):
             kwargs.pop('chat_template')
             kwargs.pop('backend')
             kwargs.pop('device')
+            kwargs.pop('eager_mode')
             kwargs['chat_template_config'] = chat_template_config
             run_chat(**kwargs)
 

lmdeploy/cli/serve.py

@@ -59,11 +59,12 @@ def add_parser_gradio():
 
         # pytorch engine args
         pt_group = parser.add_argument_group('PyTorch engine arguments')
+        ArgumentHelper.device(pt_group)
+        ArgumentHelper.eager_mode(pt_group)
 
         # common engine args
         dtype_act = ArgumentHelper.dtype(pt_group)
         tp_act = ArgumentHelper.tp(pt_group)
-        ArgumentHelper.device(pt_group)
         session_len_act = ArgumentHelper.session_len(pt_group)
         max_batch_size_act = ArgumentHelper.max_batch_size(pt_group)
         cache_max_entry_act = ArgumentHelper.cache_max_entry_count(pt_group)
@@ -159,6 +160,8 @@ def add_parser_api_server():
 
         ArgumentHelper.adapters(pt_group)
         ArgumentHelper.device(pt_group)
+        ArgumentHelper.eager_mode(pt_group)
+
         # common engine args
         dtype_act = ArgumentHelper.dtype(pt_group)
         tp_act = ArgumentHelper.tp(pt_group)
@@ -261,6 +264,7 @@ def gradio(args):
                 enable_prefix_caching=args.enable_prefix_caching,
                 device_type=args.device,
                 quant_policy=args.quant_policy,
+                eager_mode=args.eager_mode,
                 max_prefill_token_num=args.max_prefill_token_num)
         else:
             backend_config = TurbomindEngineConfig(
@@ -311,6 +315,7 @@ def api_server(args):
                 enable_prefix_caching=args.enable_prefix_caching,
                 device_type=args.device,
                 quant_policy=args.quant_policy,
+                eager_mode=args.eager_mode,
                 max_prefill_token_num=args.max_prefill_token_num)
         else:
             from lmdeploy.messages import TurbomindEngineConfig

lmdeploy/cli/utils.py

@@ -495,3 +495,13 @@ def disable_fastapi_docs(parser):
                                    default=False,
                                    help="Disable FastAPI's OpenAPI schema,"
                                    ' Swagger UI, and ReDoc endpoint')
+
+    @staticmethod
+    def eager_mode(parser):
+        """Add argument eager_mode to parser."""
+
+        return parser.add_argument('--eager-mode',
+                                   action='store_true',
+                                   default=False,
+                                   help='Whether to enable eager mode. '
+                                   'If True, cuda graph would be disabled')