Skip to content

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

Merged
terrykong merged 3 commits into
mainfrom
user/joyang/dsv3_dtensor_v2
Sep 25, 2025
Merged

terrykong merged 3 commits into
mainfrom
user/joyang/dsv3_dtensor_v2

Conversation

@joyang-nv

@joyang-nv joyang-nv commented Sep 15, 2025 •

Copy link
Copy Markdown
Member

What does this PR do ?

When you download some HF models with modules like moonlight, venvs might has assertion which doesn't have transformers_modules module. This is due to downloaded modules path is not added to PYTHONPATH .

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

Summary by CodeRabbit

  • Bug Fixes

    • Resolved intermittent failures where Transformer-related modules could be missing on distributed workers, improving model loading and tokenizer deserialization reliability.
  • Chores

    • Enhanced worker environment setup so Transformer module paths are consistently available across virtual environments.
    • No changes to public APIs.

@coderabbitai

coderabbitai Bot commented Sep 15, 2025 •

Copy link
Copy Markdown
Contributor

Walkthrough

Adds patch_transformers_module_dir to venv utilities and calls it when constructing worker runtime environments to insert the Transformers modules directory into PYTHONPATH after setting virtual-environment variables. No public API signatures changed.

Changes

Cohort / File(s) Summary
Worker env setup integration
nemo_rl/distributed/worker_groups.py
Imports patch_transformers_module_dir and calls it after setting VIRTUAL_ENV and UV_PROJECT_ENVIRONMENT, updating runtime_env["env_vars"] so the Transformers modules directory is present in PYTHONPATH before spawning workers.
Venv utilities enhancement
nemo_rl/utils/venvs.py
Adds patch_transformers_module_dir(env_vars: dict[str, str]) which locates the Transformers modules directory via transformers.utils.hub.TRANSFORMERS_CACHE, asserts the directory exists, prepends or sets it in PYTHONPATH, and returns the updated env vars.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant Caller
  participant WorkerGroups
  participant venvs as venvs.patch_transformers_module_dir
  participant Env as runtime_env["env_vars"]
  participant Worker

  Caller->>WorkerGroups: create_workers_from_bundle_indices(...)
  WorkerGroups->>Env: set VIRTUAL_ENV, UV_PROJECT_ENVIRONMENT
  WorkerGroups->>venvs: patch_transformers_module_dir(Env)
  venvs-->>WorkerGroups: Env updated (PYTHONPATH includes modules dir)
  WorkerGroups->>Worker: spawn worker with runtime_env
  Worker-->>WorkerGroups: worker ready
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

I hop through venvs, keen and spry,
I tuck modules in PYTHONPATH high.
Tokenizers wake without a fright,
Workers spring to work — all right! 🐇✨

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check ✅ Passed The title "fix: can't find transformers_modules error for moonlight" is concise and directly names the primary issue this PR addresses — a ModuleNotFoundError for transformers_modules affecting moonlight — and aligns with the changes that patch the transformers module directory into the worker environment. It uses a clear "fix:" prefix and is specific enough for a reviewer scanning history to understand the main intent without listing implementation details.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
✨ Finishing touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch user/joyang/dsv3_dtensor_v2

Tip

👮 Agentic pre-merge checks are now available in preview!

Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.

  • Built-in checks – Quickly apply ready-made checks to enforce title conventions, require pull request descriptions that follow templates, validate linked issues for compliance, and more.
  • Custom agentic checks – Define your own rules using CodeRabbit’s advanced agentic capabilities to enforce organization-specific policies and workflows. For example, you can instruct CodeRabbit’s agent to verify that API documentation is updated whenever API schema files are modified in a PR. Note: Upto 5 custom checks are currently allowed during the preview period. Pricing for this feature will be announced in a few weeks.

Please see the documentation for more information.

Example:

reviews:
  pre_merge_checks:
    custom_checks:
      - name: "Undocumented Breaking Changes"
        mode: "warning"
        instructions: |
          Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).

Please share your feedback with us on this Discord post.


Comment @coderabbitai help to get the list of available commands and usage tips.

@terrykong
terrykong requested a review from yfw September 15, 2025 15:41
@terrykong

Copy link
Copy Markdown
Collaborator

@yfw to review

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 94a3d49 and 2bd8c9e.

📒 Files selected for processing (2)
  • nemo_rl/distributed/worker_groups.py (2 hunks)
  • nemo_rl/utils/venvs.py (1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
nemo_rl/distributed/worker_groups.py (1)
nemo_rl/utils/venvs.py (2)
  • create_local_venv_on_each_node (152-189)
  • patch_transformers_module_dir (193-203)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
  • GitHub Check: CI quality check
🔇 Additional comments (1)
nemo_rl/distributed/worker_groups.py (1)

30-30: Import looks fine.

No coupling concerns; colocating venv helpers here is reasonable.

Comment thread nemo_rl/distributed/worker_groups.py
Comment thread nemo_rl/utils/venvs.py Outdated
Comment thread nemo_rl/utils/venvs.py
@terrykong terrykong changed the title Fix can't find transformers_modules error. fix: can't find transformers_modules error for moonlight Sep 17, 2025
Comment thread nemo_rl/utils/venvs.py
@github-actions

Copy link
Copy Markdown

ℹ️ File Consistency Check

Check based on commit: d32f4e6 (PR #1124 from user/joyang/dsv3_dtensor_v2)

This is a test comment


This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.

@github-actions

Copy link
Copy Markdown

ℹ️ File Consistency Check

Check based on commit: 6dda8a3 (PR #1124 from user/joyang/dsv3_dtensor_v2)

This is a test comment


This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.

@github-actions

Copy link
Copy Markdown

ℹ️ File Consistency Check

Check based on commit: 20e99a1 (PR #1124 from user/joyang/dsv3_dtensor_v2)

This is a test comment


This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.

yfw
yfw previously approved these changes Sep 17, 2025
terrykong
terrykong previously approved these changes Sep 17, 2025
@terrykong
terrykong dismissed stale reviews from yfw and themself via dac601a September 17, 2025 17:01
@terrykong
terrykong enabled auto-merge September 17, 2025 17:01
@github-actions

Copy link
Copy Markdown

ℹ️ File Consistency Check

Check based on commit: dac601a (PR #1124 from user/joyang/dsv3_dtensor_v2)

This is a test comment


This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.

terrykong
terrykong previously approved these changes Sep 17, 2025

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

♻️ Duplicate comments (2)
nemo_rl/utils/venvs.py (2)

208-213: Use os.pathsep and de-duplicate PYTHONPATH entries.

Hardcoding ":" is not cross‑platform and duplicate entries bloat envs. The diff above addresses both by using os.pathsep and checking membership before prepending.


202-207: Avoid driver-side transformers import and hard assert; compute HF modules path robustly, no-op if absent.

Importing transformers here can fail on the driver (only installed inside worker venvs) and the hard assert bricks clusters before any modules are downloaded. Prefer a resilient lookup (env -> huggingface_hub.constants -> conservative defaults), use a no-op with a debug log if missing, and add a return type. Also normalize/expand paths.

Apply this diff:

-def patch_transformers_module_dir(env_vars: dict[str, str]):
-    """Updates the PYTHONPATH in the given env_vars to include downloaded Transformers modules.
+def patch_transformers_module_dir(env_vars: dict[str, str]) -> dict[str, str]:
+    """Prepend the Hugging Face dynamic modules cache to PYTHONPATH so imports like
+    `transformers_modules.*` resolve at runtime. No-op if the cache dir is absent.

-    Some models, such as `moonshotai/Moonlight-16B-A3B`, may raise
-    `ModuleNotFoundError: No module named 'transformers_modules'` during
-    tokenizer deserialization when passed to a `Policy` object. This patch
-    prevents that error.
+    Some models (e.g., Moonlight) raise `ModuleNotFoundError: transformers_modules`
+    during tokenizer deserialization on Ray workers. This ensures the modules cache
+    is discoverable without requiring `transformers` on the driver.

-    Assumption: the cache directory is shared across virtual environments.
+    Assumption: the cache directory is shared across virtual environments.
     """
-    from transformers.utils.hub import TRANSFORMERS_CACHE
-
-    module_dir = os.path.join(TRANSFORMERS_CACHE, "..", "modules")
-    assert os.path.exists(module_dir), (
-        f"TRANSFORMERS_CACHE is not set or does not exist: {module_dir}"
-    )
-    if "PYTHONPATH" not in env_vars:
-        env_vars["PYTHONPATH"] = module_dir
-    else:
-        env_vars["PYTHONPATH"] = f"{module_dir}:{env_vars['PYTHONPATH']}"
-
-    return env_vars
+    # 1) Explicit override if provided
+    module_dir: str | None = os.environ.get("HF_MODULES_CACHE")
+    if module_dir:
+        module_dir = os.path.expanduser(module_dir)
+    else:
+        # 2) Prefer lightweight HF Hub constants (no heavy import on the driver)
+        try:
+            from huggingface_hub.constants import HF_MODULES_CACHE as _HF_MODULES_CACHE  # type: ignore
+            module_dir = os.path.expanduser(_HF_MODULES_CACHE)
+        except Exception:
+            # 3) Last-resort default (avoid importing transformers on the driver)
+            module_dir = os.path.join(os.path.expanduser("~/.cache/huggingface"), "modules")
+
+    module_dir = os.path.normpath(module_dir)
+    if not os.path.isdir(module_dir):
+        logger.debug("HF modules cache not found at %s; skipping PYTHONPATH patch.", module_dir)
+        return env_vars
+
+    existing = env_vars.get("PYTHONPATH", "")
+    parts = [p for p in existing.split(os.pathsep) if p] if existing else []
+    if module_dir not in parts:
+        env_vars["PYTHONPATH"] = os.pathsep.join([module_dir] + parts) if parts else module_dir
+    return env_vars
🧹 Nitpick comments (2)
nemo_rl/utils/venvs.py (2)

192-213: Optional: consider lightweight huggingface_hub fallback check.

If you prefer keeping a fallback that derives the path from transformers when available, we can add a nested try importing TRANSFORMERS_CACHE and compute ../modules. Omitted in the diff to avoid heavy imports on drivers.

If you want that variant, I can provide an alternative diff that adds the guarded transformers fallback.


192-201: Docstring: clarify behavior and call-site assumptions.

Add one sentence to the function docstring stating it’s a safe no‑op when the Transformers modules cache is absent and that it should be called only after runtime_env["env_vars"] is fully composed (e.g., after VIRTUAL_ENV / UV_PROJECT_ENVIRONMENT are set). Call site confirmed: nemo_rl/distributed/worker_groups.py invokes patch_transformers_module_dir(runtime_env["env_vars"]) after setting those env vars (line 516).

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 20e99a1 and dac601a.

📒 Files selected for processing (2)
  • nemo_rl/distributed/worker_groups.py (2 hunks)
  • nemo_rl/utils/venvs.py (1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • nemo_rl/distributed/worker_groups.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
  • GitHub Check: Lint check
  • GitHub Check: Post submodule check comment / Comment on PR
🔇 Additional comments (1)
nemo_rl/utils/venvs.py (1)

192-213: Ignore HF_MODULES_CACHE — it's not a public constant; keep using TRANSFORMERS_CACHE.

huggingface_hub.constants does not export HF_MODULES_CACHE; the current approach (from transformers.utils.hub import TRANSFORMERS_CACHE and deriving the modules dir via TRANSFORMERS_CACHE/../modules) is correct.

Likely an incorrect or invalid review comment.

@terrykong
terrykong added this pull request to the merge queue Sep 17, 2025
@chtruong814
chtruong814 removed this pull request from the merge queue due to the queue being cleared Sep 18, 2025
@chtruong814
chtruong814 added this pull request to the merge queue Sep 18, 2025
@chtruong814
chtruong814 removed this pull request from the merge queue due to the queue being cleared Sep 18, 2025
@chtruong814
chtruong814 added this pull request to the merge queue Sep 18, 2025
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Sep 18, 2025
@terrykong terrykong added the CI:L1 Run doctests, unit tests, and functional tests label Sep 18, 2025
@terrykong
terrykong enabled auto-merge (squash) September 18, 2025 18:20
@chtruong814

chtruong814 commented Sep 18, 2025 •

Copy link
Copy Markdown
Contributor

@joyang-nv @terrykong I think this PR has an issue. When CI ran in the queue, it failed with this in the functional tests

https://github.com/NVIDIA-NeMo/RL/actions/runs/17822565027/job/50687190084

  File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 397, in __init__
    self._create_workers_from_bundle_indices(
  File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 516, in _create_workers_from_bundle_indices
    patch_transformers_module_dir(runtime_env["env_vars"])
  File "/opt/nemo-rl/nemo_rl/utils/venvs.py", line 205, in patch_transformers_module_dir
    assert os.path.exists(module_dir), (
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: TRANSFORMERS_CACHE is not set or does not exist: /home/TestData/nemo-rl/hf_home/hub/../modules
Loaded configuration from: /opt/nemo-rl/examples/configs/sft.yaml
Overrides: ['policy.model_name=Qwen/Qwen3-0.6B', 'cluster.gpus_per_node=2', 'sft.max_num_steps=3', 'sft.val_batches=1', 'sft.val_period=3', 'logger.tensorboard_enabled=true', 'logger.log_dir=/opt/nemo-rl/tests/functional/sft/logs', 'logger.wandb_enabled=false', 'logger.monitor_gpus=true', 'checkpointing.enabled=true', 'checkpointing.save_period=3', 'checkpointing.checkpoint_dir=/tmp/sft_checkpoints']
Applied CLI overrides

I see. Thanks for info. I will fix it.

@joyang-nv
joyang-nv requested review from a team as code owners September 24, 2025 15:45
@joyang-nv joyang-nv added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Sep 24, 2025
@joyang-nv joyang-nv removed the CI:L1 Run doctests, unit tests, and functional tests label Sep 24, 2025
@joyang-nv joyang-nv added the CI:L1 Run doctests, unit tests, and functional tests label Sep 24, 2025
@joyang-nv
joyang-nv requested a review from terrykong September 24, 2025 15:54
@joyang-nv joyang-nv removed the CI:L1 Run doctests, unit tests, and functional tests label Sep 24, 2025
Signed-off-by: Jonas Yang <joyang@nvidia.com>
@joyang-nv
joyang-nv force-pushed the user/joyang/dsv3_dtensor_v2 branch from 8b3437b to 3fef216 Compare September 24, 2025 16:01
@joyang-nv joyang-nv added the CI:L1 Run doctests, unit tests, and functional tests label Sep 24, 2025
@joyang-nv

Copy link
Copy Markdown
Member Author

@terrykong can you approve this once you saw CI pass? And not sure whether still need @chtruong814 to help merge?

@terrykong terrykong added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Sep 24, 2025
@terrykong
terrykong merged commit a579137 into main Sep 25, 2025
25 checks passed
@terrykong
terrykong deleted the user/joyang/dsv3_dtensor_v2 branch September 25, 2025 01:30
jseppanen pushed a commit that referenced this pull request Oct 9, 2025
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>
PrinsYin pushed a commit to PrinsYin/RL that referenced this pull request Nov 30, 2025
…#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>
yuanhangsu1986 pushed a commit to yuanhangsu1986/RL-Nemontron-Edge-Omni that referenced this pull request Feb 21, 2026
…#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>
Signed-off-by: yuanhangs <yuanhangs@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L1 Run doctests, unit tests, and functional tests r0.4.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants