[https://nvbugs/5556998][fix] init_hf_modules in worker_main for models with trust_remote=true (NVIDIA#8931)

Signed-off-by: Lanyu Liao <[email protected]>
Co-authored-by: Lanyu Liao <[email protected]>

Author: lancelly

tensorrt_llm/executor/base_worker.py

@@ -40,9 +40,32 @@
 
 __all__ = [
     "BaseWorker",
+    "_init_hf_modules",
 ]
 
 
+def _init_hf_modules():
+    """Initialize cached HuggingFace modules for models with trust_remote_code=True.
+
+    This is safe to call multiple times (idempotent) and should be called:
+    1. At module import time (for main process and spawned subprocesses)
+    2. At worker_main entry (for forked processes or external MPI ranks)
+
+    References: https://github.com/vllm-project/vllm/pull/871
+    """
+    try:
+        from transformers.dynamic_module_utils import init_hf_modules
+        init_hf_modules()
+        logger.debug("HF modules initialized")
+    except ImportError as e:
+        logger.warning(f"ImportError initializing HF modules: {e}")
+    except Exception as e:
+        logger.error(f"Exception initializing HF modules: {e}")
+
+
+_init_hf_modules()
+
+
 class BaseWorker(GenerationExecutor):
 
     class WorkerExit(GeneratorExit):

tensorrt_llm/executor/worker.py

@@ -21,7 +21,7 @@
 from ..llmapi.utils import (AsyncQueue, ManagedThread, _SyncQueue, logger_debug,
                             print_traceback_on_error)
 from ..sampling_params import BatchedLogitsProcessor
-from .base_worker import BaseWorker
+from .base_worker import BaseWorker, _init_hf_modules
 from .executor import IterationResultQueue
 from .ipc import FusedIpcQueue, IpcQueue
 from .postproc_worker import (PostprocWorker, PostprocWorkerConfig,
@@ -242,6 +242,10 @@ def worker_main(
     llm_args: Optional[BaseLlmArgs] = None,
 ) -> None:
     mpi_comm().barrier()
+
+    if llm_args is not None and llm_args.trust_remote_code:
+        _init_hf_modules()
+
     logger_debug(f"Worker {mpi_rank()} entering worker_main...\n", "green")
 
     result_queue: Optional[IpcQueue] = None