No module named 'nemo'

[tpoisonooo]

Describe the bug

When using megatron, ray node outputs

  File "/data/khj/workspace/RL/nemo_rl/models/policy/megatron_policy_worker.py", line 60, in <module>
    from nemo.tron import fault_tolerance
ModuleNotFoundError: No module named 'nemo'

Steps/Code to reproduce bug

Step1

uv venv
bash tools/build-flash-attn-in-uv-cache.sh

Step2

Since github is unreachable, manully download flash-attn, update pyproject.toml in nvidia-nemo/RL. Then uv sync --extra vllm and uv sync --extra mcore

Here is all my update, no new feature.

Step3. Run test and passed

uv run python examples/run_grpo_math.py

Step4. Try megatron, but crash

1. update run_grpo_math.py, use config/grpo_math_1B_megatron.yaml
2. rerun uv run python examples/run_grpo_math.py

  File "/data/khj/workspace/RL/.venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2822, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/khj/workspace/RL/.venv/lib/python3.12/site-packages/ray/_private/worker.py", line 930, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(ModuleNotFoundError): ray::IsolatedWorkerInitializer.create_worker() (pid=1125989, ip=10.15.1.122, actor_id=426525ba1cb791498a2fa12001000000, repr=<nemo_rl.distributed.worker_groups.IsolatedWorkerInitializer object at 0x7fd0ddd0b740>)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/khj/workspace/RL/nemo_rl/distributed/worker_groups.py", line 165, in create_worker
    module = importlib.import_module(module_name)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/khj/miniconda3/lib/python3.12/importlib/__init__.py", line 90, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap>", line 1387, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1360, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1331, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 935, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 999, in exec_module
  File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
  File "/data/khj/workspace/RL/nemo_rl/models/policy/megatron_policy_worker.py", line 60, in <module>
    from nemo.tron import fault_tolerance
ModuleNotFoundError: No module named 'nemo'

Some trial

1. print env in megatron_policy_worker.py

ray node is not using RL/.env/bin/python, use RL/venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/bin/python instead:

(IsolatedWorkerInitializer pid=1125989) !!!Ray Python executable: /data/khj/workspace/RL/venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/bin/python
(IsolatedWorkerInitializer pid=1125989) !!!Ray sys.path: ['/data/khj/workspace/RL', '/data/khj/workspace/RL/examples', '/data/khj/workspace/RL/venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/lib/python3.12/site-packages/ray/thirdparty_files', '/data/khj/workspace/RL/.venv/lib/python3.12/site-packages/ray/_private/workers', '/home/khj/miniconda3/lib/python312.zip', '/home/khj/miniconda3/lib/python3.12', '/home/khj/miniconda3/lib/python3.12/lib-dynload', '/data/khj/workspace/RL/venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/lib/python3.12/site-packages', '/data/khj/workspace/RL/3rdparty/Megatron-LM-workspace/Megatron-LM', '/tmp/tmpiueokp9q', '/data/khj/workspace/RL/venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/lib/python3.12/site-packages/setuptools/_vendor']

So, I try to specify env in ray_action_environment_registry.py, but it still crash.

Environment overview (please complete the following information)

• Environment location: Bare-metal
• Method of install: from source
• not docker, I install it on host.

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• OS version: centOS
• PyTorch version: 2.7.1
• Python version: 3.12.11

Additional context

Add any other context about the problem here.
GPU H800

[terrykong]

Hi @tpoisonooo . Thanks for the interest in NeMo RL

When you setup, did you download the submodules in the prerequisites

git submodule update --init --recursive

Assuming the submodule is checked out under 3rdparty/NeMo-workspace/NeMo, another thing to try is to force the venvs to build like in these instructions. What can happen is if someone wasn't using the megatron path and runs some dtensor experiments, but decides later they want to try the mcore path, the venvs have to be rebuilt since we create a "stub" package for megatron to fake its installation so you need to force it to rebuild.

[tpoisonooo]

Great thanks @terrykong , after more modifications, it works !

Modification 1

I mistakenly downloaded nvidia/Megatron-LM instead of terrykong/Megatron-LM.

Since

• Github/huggingface unreacheable here (bad GFW)
• git.proxy does not work for this kind of proxy: http://ghfast.top/https://github.com/NVIDIA-NeMo/RL

Hopes that RL/uv.pyproject also supports github/HF proxy.

Modification 2

Hardcode megatron_checkpoint_home to my local dir.

[terrykong]

Glad you got it up and running!

Hopes that RL/uv.pyproject also supports github/HF proxy.

Do you mind opening up a feature request against uv? Feel free to link to this issue

Hardcode megatron_checkpoint_home to my local dir.

I think this will probably help #665

[terrykong]

@tpoisonooo do you still require further support for the initial issue you raised?

[tpoisonooo]

@tpoisonooo do you still require further support for the initial issue you raised?

Thanks for your help, I already fixed the initial bug and then try to contribute some PRs.

[tpoisonooo]

Another question:

For better GPU efficiency, should I always use megatron or dtensor?

I found that sft_openmathinstruct2.yaml's’ default option is dtensor.

After read dtensor_policy_worker.py, I guess dtensor may slower than megatron.
So I use megatron in sft_openmathinstruct2.yaml,

Time comparision:

• before 15~18s per step
• after 5s per step

And it crashed on model_utils.py: backward(), I am working on it.

  File "/fs-computility/ai4agr/konghuanjun/RL/venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/lib/p
ython3.12/site-packages/torch/autograd/function.py", line 307, in apply
    return user_fn(self, *args)
           ^^^^^^^^^^^^^^^^^^^^
  File "/fs-computility/ai4agr/konghuanjun/RL/nemo_rl/distributed/model_utils.py", line 112, in backward
    is_chosen = (~target_mask).unsqueeze(-1) * torch.nn.functional.one_hot(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 11.92 GiB. GPU 0 has a total capacity of 79.33 GiB of which 1
0.37 GiB is free. Process 2887511 has 68.95 GiB memory in use. Of the allocated memory 64.02 GiB is allocated by PyTorch, a
nd 1.33 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_AL
LOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/do
cs/stable/notes/cuda.html#environment-variables)

[tpoisonooo]

Another question:

For better GPU efficiency, should I always use megatron or dtensor?

I found that sft_openmathinstruct2.yaml's’ default option is dtensor.

After read dtensor_policy_worker.py, I guess dtensor may slower than megatron. So I use megatron in sft_openmathinstruct2.yaml,

Time comparision:

• before 15~18s per step
• after 5s per step

And it crashed on model_utils.py: backward(), I am working on it.

File "/fs-computility/ai4agr/konghuanjun/RL/venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/lib/p
ython3.12/site-packages/torch/autograd/function.py", line 307, in apply
return user_fn(self, *args)
^^^^^^^^^^^^^^^^^^^^
File "/fs-computility/ai4agr/konghuanjun/RL/nemo_rl/distributed/model_utils.py", line 112, in backward
is_chosen = (~target_mask).unsqueeze(-1) * torch.nn.functional.one_hot(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 11.92 GiB. GPU 0 has a total capacity of 79.33 GiB of which 1
0.37 GiB is free. Process 2887511 has 68.95 GiB memory in use. Of the allocated memory 64.02 GiB is allocated by PyTorch, a
nd 1.33 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_AL
LOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/do
cs/stable/notes/cuda.html#environment-variables)

I think I have optimized the implementation of backward(), a PR would be send later.

[terrykong]

@tpoisonooo shared this PR #761 and @SahilJain314 helped merge it with tests (giving you credit for the contribution @tpoisonooo) #761

Since the original issue has been solved, I'll close this. Feel free to open new issues as you encounter. Thanks for making nemo rl better!