fix autodeploy tokenizer

Signed-off-by: leslie-fang25 <[email protected]>

Author: leslie-fang25

tensorrt_llm/executor/executor.py

@@ -25,6 +25,7 @@
 from ..llmapi.llm_utils import KvCacheRetentionConfig
 from ..llmapi.mpi_session import (MpiSession, external_mpi_comm_available,
                                   need_spawn_mpi_workers)
+from ..llmapi.tokenizer import TokenizerBase
 from ..llmapi.utils import (AsyncQueue, enable_llm_debug,
                             enable_worker_single_process_for_tp1, print_colored,
                             print_colored_debug)
@@ -356,6 +357,7 @@ def create(
         is_llm_executor: Optional[bool] = None,
         lora_config: Optional[LoraConfig] = None,
         hf_model_dir: Optional[Path] = None,
+        tokenizer: Optional[TokenizerBase] = None,
         llm_args: Optional[TorchLlmArgs] = None,
     ) -> Union["GenerationExecutorProxy", "GenerationExecutorWorker"]:
         # local imports to avoid cyclic importing
@@ -384,6 +386,7 @@ def create(
             "executor_config": executor_config,
             "batched_logits_processor": batched_logits_processor,
             "hf_model_dir": hf_model_dir,
+            "tokenizer": tokenizer,
             "llm_args": llm_args,
         }
 

tensorrt_llm/executor/worker.py

@@ -20,6 +20,7 @@
 from ..builder import ConfigEncoder, Engine, EngineConfig
 from ..llmapi.llm_args import PybindMirror, TorchLlmArgs
 from ..llmapi.mpi_session import set_mpi_session_cpp
+from ..llmapi.tokenizer import TokenizerBase
 from ..llmapi.tracer import VizTracer, global_tracer, set_global_tracer
 from ..llmapi.utils import (AsyncQueue, ManagedThread, _SyncQueue,
                             clear_sched_affinity, print_colored_debug,
@@ -61,6 +62,7 @@ def __init__(
         is_llm_executor: Optional[bool] = None,
         lora_config: Optional[LoraConfig] = None,
         hf_model_dir: Optional[Path] = None,
+        tokenizer: Optional[TokenizerBase] = None,
         llm_args: Optional[TorchLlmArgs] = None,
     ) -> None:
         postproc_config = postproc_worker_config or PostprocWorkerConfig()
@@ -102,7 +104,8 @@ def _get_comm_ranks_device_id():
 
         def _create_py_executor(executor_config):
             assert executor_config is None, "expect an empty executor_config is _create_py_executor"
-            executor_config = llm_args.get_executor_config(hf_model_dir)
+            executor_config = llm_args.get_executor_config(
+                hf_model_dir, tokenizer)
             # Persist so downstream code (e.g., default max_tokens deduction) has access
             self._executor_config = executor_config
             executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig(
@@ -662,6 +665,7 @@ def worker_main(
         bool] = True,  # whether it's the main executor instance
     lora_config: Optional[LoraConfig] = None,
     hf_model_dir: Optional[Path] = None,
+    tokenizer: Optional[TokenizerBase] = None,
     llm_args: Optional[TorchLlmArgs] = None,
 ) -> None:
     mpi_comm().barrier()
@@ -790,6 +794,7 @@ def notify_proxy_threads_to_quit():
             is_llm_executor=is_llm_executor,
             lora_config=lora_config,
             hf_model_dir=hf_model_dir,
+            tokenizer=tokenizer,
             llm_args=llm_args)
     except Exception as e:
         logger.error(f"Failed to initialize executor on rank {mpi_rank()}: {e}")

tensorrt_llm/llmapi/llm.py

@@ -966,9 +966,6 @@ def _build_model(self):
                                                       self.tokenizer)
         self._tokenizer = self.input_processor.tokenizer
 
-        # Update the tokenizer in TorchLlmArgs, so it can be used in GenerationExecutorWorker to init executor_config
-        self.args.set_tokenizer(self.tokenizer)
-
         # TODO: revisit gather_context_logits
         return_logits = self.args.gather_generation_logits
         self._executor = self._executor_cls.create(
@@ -987,6 +984,7 @@ def _build_model(self):
             is_llm_executor=True,
             lora_config=self.args.lora_config,
             hf_model_dir=self._hf_model_dir,
+            tokenizer=self.tokenizer,
             llm_args=self.args)
 
     def _validate_args_for_torch_backend(self, kwargs: dict) -> None:

tensorrt_llm/llmapi/llm_args.py

@@ -1837,12 +1837,11 @@ def _load_config_from_ckpt(self, ckpt_dir: Path):
                 moe_tp_size=moe_tp_size,
                 moe_ep_size=moe_ep_size)
 
-    def set_tokenizer(self, tokenizer):
-        self.tokenizer = tokenizer
-
-    def get_executor_config(self,
-                            _hf_model_dir: Optional[Path] = None
-                            ) -> _ExecutorConfig:
+    def get_executor_config(
+        self,
+        _hf_model_dir: Optional[Path] = None,
+        tokenizer: Optional[TokenizerBase] = None,
+    ) -> _ExecutorConfig:
         executor_config = _ExecutorConfig(
             max_beam_width=self.max_beam_width,
             scheduler_config=PybindMirror.maybe_to_pybind(
@@ -1867,13 +1866,15 @@ def get_executor_config(self,
         if self.decoding_config is not None:
             executor_config.decoding_config = self.decoding_config
         if self.guided_decoding_backend == 'xgrammar':
+            assert tokenizer is not None
             executor_config.guided_decoding_config = _GuidedDecodingConfig(
                 backend=_GuidedDecodingConfig.GuidedDecodingBackend.XGRAMMAR,
-                **_xgrammar_tokenizer_info(self.tokenizer))
+                **_xgrammar_tokenizer_info(tokenizer))
         elif self.guided_decoding_backend == 'llguidance':
+            assert tokenizer is not None
             executor_config.guided_decoding_config = _GuidedDecodingConfig(
                 backend=_GuidedDecodingConfig.GuidedDecodingBackend.LLGUIDANCE,
-                **_llguidance_tokenizer_info(self.tokenizer))
+                **_llguidance_tokenizer_info(tokenizer))
         elif self.guided_decoding_backend is not None:
             raise ValueError(
                 f"Unsupported guided decoding backend {self.guided_decoding_backend}"
@@ -2460,10 +2461,12 @@ def validate_batch_wait_timeout_ms(self) -> 'TorchLlmArgs':
     def set_mm_encoder_only(self, mm_encoder_only):
         self.mm_encoder_only = mm_encoder_only
 
-    def get_executor_config(self,
-                            _hf_model_dir: Optional[Path] = None
-                            ) -> _ExecutorConfig:
-        executor_config = super().get_executor_config(_hf_model_dir)
+    def get_executor_config(
+        self,
+        _hf_model_dir: Optional[Path] = None,
+        tokenizer: Optional[TokenizerBase] = None,
+    ) -> _ExecutorConfig:
+        executor_config = super().get_executor_config(_hf_model_dir, tokenizer)
         executor_config.mm_encoder_only = self.mm_encoder_only
         return executor_config
 

tensorrt_llm/llmapi/mm_encoder.py

@@ -56,8 +56,6 @@ def _build_model(self):
         self._tokenizer = self.input_processor.tokenizer
 
         assert isinstance(self.args, TorchLlmArgs)
-        # Update the tokenizer in TorchLlmArgs, so it can be used in GenerationExecutorWorker to init executor_config
-        self.args.set_tokenizer(self.tokenizer)
         self.args.set_mm_encoder_only(True)
 
         self._executor = self._executor_cls.create(
@@ -69,6 +67,7 @@ def _build_model(self):
                 self.args.parallel_config.world_size),
             is_llm_executor=True,  # TODO: check if this is correct or needed
             hf_model_dir=self._hf_model_dir,
+            tokenizer=self.tokenizer,
             llm_args=self.args)
 
     def _validate_mm_args_for_torch_backend(self, kwargs: dict) -> None: