Add inference_max_seq_len to ray mbridge deployment path

nemo_deploy/deploy_ray.py

@@ -187,6 +187,7 @@ def deploy_inframework_model(
         enable_flash_decode: bool = False,
         legacy_ckpt: bool = False,
         max_batch_size: int = 32,
+        inference_max_seq_length: int = 4096,
         random_seed: Optional[int] = None,
         test_mode: bool = False,
         model_type: str = "gpt",
@@ -254,6 +255,7 @@ def deploy_inframework_model(
                 enable_flash_decode=enable_flash_decode,
                 legacy_ckpt=legacy_ckpt,
                 max_batch_size=max_batch_size,
+                inference_max_seq_length=inference_max_seq_length,
                 random_seed=random_seed,
                 model_type=model_type,
                 micro_batch_size=micro_batch_size,

nemo_deploy/llm/megatronllm_deployable_ray.py

@@ -58,6 +58,7 @@ def __init__(
         enable_flash_decode: bool = False,
         legacy_ckpt: bool = False,
         max_batch_size: int = 32,
+        inference_max_seq_length: int = 4096,
         random_seed: Optional[int] = None,
         model_type: str = "gpt",
         micro_batch_size: Optional[int] = None,
@@ -93,6 +94,7 @@ def __init__(
                 enable_flash_decode=enable_flash_decode,
                 legacy_ckpt=legacy_ckpt,
                 max_batch_size=max_batch_size,
+                inference_max_seq_length=inference_max_seq_length,
                 random_seed=random_seed,
                 model_type=model_type,
                 micro_batch_size=micro_batch_size,
@@ -135,6 +137,7 @@ def __init__(
         enable_flash_decode: bool = False,
         legacy_ckpt: bool = False,
         max_batch_size: int = 32,
+        inference_max_seq_length: int = 4096,
         random_seed: Optional[int] = None,
         model_type: str = "gpt",
         micro_batch_size: Optional[int] = None,
@@ -152,12 +155,10 @@ def __init__(
             enable_cuda_graphs (bool): Whether to enable CUDA graphs for faster inference.
             enable_flash_decode (bool): Whether to enable Flash Attention decode.
             max_batch_size (int): Maximum batch size for request batching.
-            batch_wait_timeout_s (float): Maximum time to wait for batching requests.
+            inference_max_seq_length (int): Maximum sequence length for inference.
             legacy_ckpt (bool): Whether to use legacy checkpoint format. Defaults to False.
             random_seed (int): Random seed for model initialization.
-            megatron_checkpoint_filepath (str): Path to the Megatron checkpoint file.
             model_type (str): Type of model to load.
-            model_format (str): Format of model to load.
             micro_batch_size (Optional[int]): Micro batch size for model execution.
         """
         try:
@@ -202,6 +203,7 @@ def __init__(
                 enable_flash_decode=enable_flash_decode,
                 legacy_ckpt=legacy_ckpt,
                 max_batch_size=max_batch_size,
+                inference_max_seq_length=inference_max_seq_length,
                 random_seed=random_seed,
                 model_type=model_type,
                 micro_batch_size=micro_batch_size,
@@ -238,6 +240,7 @@ def __init__(
                     enable_cuda_graphs=enable_cuda_graphs,
                     enable_flash_decode=enable_flash_decode,
                     max_batch_size=max_batch_size,
+                    inference_max_seq_length=inference_max_seq_length,
                     random_seed=random_seed,
                     model_type=model_type,
                     micro_batch_size=micro_batch_size,

scripts/deploy/llm/mbridge/deploy_ray.py

@@ -161,6 +161,12 @@ def parse_args():
         default=32,
         help="Maximum batch size for inference",
     )
+    parser.add_argument(
+        "--inference_max_seq_length",
+        type=int,
+        default=4096,
+        help="Maximum sequence length for inference",
+    )
     parser.add_argument(
         "--random_seed",
         type=int,
@@ -227,6 +233,7 @@ def main():
         enable_flash_decode=args.enable_flash_decode,
         legacy_ckpt=args.legacy_ckpt,
         max_batch_size=args.max_batch_size,
+        inference_max_seq_length=args.inference_max_seq_length,
         random_seed=args.random_seed,
         model_type=args.model_type,
         micro_batch_size=args.micro_batch_size,

scripts/deploy/nlp/deploy_ray_inframework.py

@@ -164,6 +164,13 @@ def parse_args():
         default=32,
         help="Maximum batch size for inference",
     )
+    parser.add_argument(
+        "-imsl",
+        "--inference_max_seq_length",
+        default=4096,
+        type=int,
+        help="Max sequence length for inference",
+    )
     parser.add_argument(
         "--random_seed",
         type=int,
@@ -244,6 +251,7 @@ def main():
         enable_flash_decode=args.enable_flash_decode,
         legacy_ckpt=args.legacy_ckpt,
         max_batch_size=args.max_batch_size,
+        inference_max_seq_length=args.inference_max_seq_length,
         random_seed=args.random_seed,
         model_type=args.model_type,
         micro_batch_size=args.micro_batch_size,