diff --git a/nemo_deploy/llm/inference/inference_base.py b/nemo_deploy/llm/inference/inference_base.py index b2f16d4fd..30aeab7ce 100644 --- a/nemo_deploy/llm/inference/inference_base.py +++ b/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: diff --git a/nemo_deploy/llm/megatronllm_deployable.py b/nemo_deploy/llm/megatronllm_deployable.py index 10910ac90..5791ab335 100755 --- a/nemo_deploy/llm/megatronllm_deployable.py +++ b/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 diff --git a/scripts/deploy/llm/mlm/deploy_triton.py b/scripts/deploy/llm/mlm/deploy_triton.py index 3fd0cdcda..51a347c1a 100755 --- a/scripts/deploy/llm/mlm/deploy_triton.py +++ b/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, ) diff --git a/scripts/deploy/nlp/deploy_inframework_triton.py b/scripts/deploy/nlp/deploy_inframework_triton.py index 3db1a46da..1f86d0143 100755 --- a/scripts/deploy/nlp/deploy_inframework_triton.py +++ b/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, ) diff --git a/scripts/deploy/nlp/deploy_ray_inframework.py b/scripts/deploy/nlp/deploy_ray_inframework.py index 23d701f26..cbabbbc26 100644 --- a/scripts/deploy/nlp/deploy_ray_inframework.py +++ b/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, )