Browse Source

fix: ModuleNotFoundError for remote code models

AlpinDale 1 year ago
parent
commit
d71a84b780
2 changed files with 5 additions and 2 deletions
  1. 1 1
      aphrodite/engine/aphrodite_engine.py
  2. 4 1
      aphrodite/engine/ray_tools.py

+ 1 - 1
aphrodite/engine/aphrodite_engine.py

@@ -153,7 +153,7 @@ class AphroditeEngine:
                     placement_group=placement_group,
                     placement_group_capture_child_tasks=True),
                 **ray_remote_kwargs,
-            )(RayWorker).remote()
+            )(RayWorker).remote(self.model_config.trust_remote_code)
             self.workers.append(worker)
 
         # Initialize torch distributed process group for the workers.

+ 4 - 1
aphrodite/engine/ray_tools.py

@@ -10,7 +10,10 @@ try:
     """Ray wrapper for aphrodite.task_handler.worker, allowing
     worker to be lazily initialized after Ray sets CUDA_VISIBLE_DEVICES."""
     class RayWorker(TorchDistributedWorker):
-        def __init__(self) -> None:
+        def __init__(self, init_cached_hf_modules=False) -> None:
+            if init_cached_hf_modules:
+                from transformers.dynamic_module_utils import init_hf_modules
+                init_hf_modules()
             self.worker = None
         
         def init_worker(self, worker_init_fn):