fix: can't find transformers_modules error for moonlight (#1124)

Signed-off-by: Jonas Yang <joyang@nvidia.com>
Signed-off-by: Jonas Yang CN <joyang@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>

Author: joyang-nv

nemo_rl/distributed/worker_groups.py

@@ -27,7 +27,10 @@
 )
 from nemo_rl.distributed.virtual_cluster import RayVirtualCluster
 from nemo_rl.distributed.worker_group_utils import recursive_merge_options
-from nemo_rl.utils.venvs import create_local_venv_on_each_node
+from nemo_rl.utils.venvs import (
+    create_local_venv_on_each_node,
+    patch_transformers_module_dir,
+)
 
 
 @dataclass
@@ -529,6 +532,7 @@ def _create_workers_from_bundle_indices(
                 }
                 runtime_env["env_vars"]["VIRTUAL_ENV"] = py_executable
                 runtime_env["env_vars"]["UV_PROJECT_ENVIRONMENT"] = py_executable
+                patch_transformers_module_dir(runtime_env["env_vars"])
 
                 extra_options = {"runtime_env": runtime_env, "name": name}
 

nemo_rl/utils/venvs.py

@@ -187,3 +187,18 @@ def create_local_venv_on_each_node(py_executable: str, venv_name: str):
     ray.util.remove_placement_group(pg)
     # Return mapping from node IP to venv python path
     return paths[0]
+
+
+# Need to set PYTHONPATH to include transformers downloaded modules.
+# Assuming the cache directory is the same cross venvs.
+def patch_transformers_module_dir(env_vars: dict[str, str]):
+    from transformers.utils.hub import TRANSFORMERS_CACHE
+
+    module_dir = os.path.join(TRANSFORMERS_CACHE, "..", "modules")
+    assert module_dir is not None, "TRANSFORMERS_CACHE should exist."
+    if "PYTHONPATH" not in env_vars:
+        env_vars["PYTHONPATH"] = module_dir
+    else:
+        env_vars["PYTHONPATH"] = f"{module_dir}:{env_vars['PYTHONPATH']}"
+
+    return env_vars