Fix MBridge tokenizer issue

nemo_deploy/llm/inference/inference_base.py

@@ -194,6 +194,7 @@ def setup_megatron_model_and_tokenizer_for_inference(
     sequence_parallel: Optional[bool] = None,
     micro_batch_size: Optional[int] = None,
     model_type: str = "gpt",
+    trust_remote_code: bool = False,
 ) -> Tuple[List[MegatronModule], MegatronTokenizer]:
     """Initialize a Megatron model and tokenizer for inference from a Megatron-LM/MBridge checkpoint.
 
@@ -259,7 +260,7 @@ def setup_megatron_model_and_tokenizer_for_inference(
         megatron_args=mlm_args,
         use_cpu_init=False,
     )
-    tokenizer = load_tokenizer(checkpoint_path)
+    tokenizer = load_tokenizer(checkpoint_path, trust_remote_code=trust_remote_code)
     return model, tokenizer, mlm_args
 
 
@@ -436,6 +437,7 @@ def create_mcore_engine(
     micro_batch_size: Optional[int] = None,
     buffer_size_gb: float = 10.0,
     legacy_model_format: bool = False,
+    trust_remote_code: bool = False,
     **model_config_kwargs,
 ) -> Tuple[MCoreEngineWithCleanup, Union[MCoreTokenizerWrappper, MegatronTokenizer]]:
     """Set up the model, tokenizer and MegatronLLM engine for inference.
@@ -497,6 +499,7 @@ def create_mcore_engine(
             sequence_parallel=sequence_parallel,
             micro_batch_size=micro_batch_size,
             model_type=model_type,
+            trust_remote_code=trust_remote_code,
         )
         model = modelList[0]
     else:

nemo_deploy/llm/megatronllm_deployable.py

@@ -105,6 +105,7 @@ def __init__(
         micro_batch_size: Optional[int] = None,
         buffer_size_gb: float = 10.0,
         legacy_model_format: bool = False,
+        trust_remote_code: bool = False,
         **model_config_kwargs,
     ):
         if not HAVE_TRITON:
@@ -136,6 +137,7 @@ def __init__(
             micro_batch_size=micro_batch_size,
             buffer_size_gb=buffer_size_gb,
             legacy_model_format=legacy_model_format,
+            trust_remote_code=trust_remote_code,
             **model_config_kwargs,
         )
         self.enable_cuda_graphs = enable_cuda_graphs

scripts/deploy/llm/mlm/deploy_triton.py

@@ -216,6 +216,13 @@ def get_args(argv):
         default=None,
         help="Micro batch size for model execution",
     )
+    parser.add_argument(
+        "-trc",
+        "--trust_remote_code",
+        default=False,
+        action="store_true",
+        help="Pass trust_remote_code=True to load_tokenizer() for checkpoints whose tokenizer requires it.",
+    )
     args = parser.parse_args(argv)
     return args
 
@@ -263,6 +270,7 @@ def nemo_deploy(argv):
         legacy_ckpt=args.legacy_ckpt,
         model_type=args.model_type,
         micro_batch_size=args.micro_batch_size,
+        trust_remote_code=args.trust_remote_code,
         **model_config_kwargs,
     )
 

scripts/deploy/nlp/deploy_inframework_triton.py

@@ -234,6 +234,13 @@ def get_args(argv):
         action="store_true",
         help="Use the legacy StaticInferenceEngine path in MCoreEngine",
     )
+    parser.add_argument(
+        "-trc",
+        "--trust_remote_code",
+        default=False,
+        action="store_true",
+        help="Pass trust_remote_code=True to load_tokenizer() for checkpoints whose tokenizer requires it.",
+    )
     args = parser.parse_args(argv)
     return args
 
@@ -284,6 +291,7 @@ def nemo_deploy(argv):
         model_type=args.model_type,
         micro_batch_size=args.micro_batch_size,
         legacy_model_format=args.legacy_model_format,
+        trust_remote_code=args.trust_remote_code,
         **model_config_kwargs,
     )
 

scripts/deploy/nlp/deploy_ray_inframework.py

@@ -225,6 +225,13 @@ def parse_args():
         default={},
         help="Runtime environment for the deployment (JSON string)",
     )
+    parser.add_argument(
+        "-trc",
+        "--trust_remote_code",
+        default=False,
+        action="store_true",
+        help="Pass trust_remote_code=True to load_tokenizer() for checkpoints whose tokenizer requires it.",
+    )
     return parser.parse_args()
 
 
@@ -283,6 +290,7 @@ def main():
         micro_batch_size=args.micro_batch_size,
         batch_wait_timeout_s=args.batch_wait_timeout_s,
         legacy_model_format=args.legacy_model_format,
+        trust_remote_code=args.trust_remote_code,
         **model_config_kwargs,
     )