fix: remove mlm workspace by kajalj22 · Pull Request #2136 · NVIDIA-NeMo/RL

kajalj22 · 2026-03-21T22:21:07Z

What does this PR do?

Removes the 3rdparty/Megatron-LM-workspace/ directory, which was a proxy workspace that re-packaged megatron-core with all its dev dependencies bundled into a flat install_requires. Instead, megatron-core is now sourced directly from the MBridge submodule at 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM.

Why

The MLM workspace was a maintenance burden — it had a hand-maintained CACHED_DEPENDENCIES list in setup.py that had to stay in sync with megatron-core's pyproject.toml. Removing it simplifies the dependency tree and eliminates a source of drift.

What changed

Deleted 3rdparty/Megatron-LM-workspace/ (setup.py, pyproject.toml, is_megatron_installed.py, submodule pointer)
Removed the .gitmodules entry and workspace member for the old submodule
Redirected megatron-core source in pyproject.toml to MBridge's Megatron-LM path
Updated the sys.path hack in nemo_rl/__init__.py to point to the new Megatron-LM location (for megatron.training importability)
Added onnxscript and nvidia-modelopt[torch] to the mcore extra — these were previously pulled in as transitive deps of the old workspace's fat install_requires and are required at module load time by transformer_engine and megatron-bridge respectively
Added nvidia-resiliency-ext git source override (needed for cp313 wheel resolution)
Updated 8+ files with stale path references (infra, tests, docs, tools)
Inlined the megatron install check in the functional test (the old is_megatron_installed module was a py_modules entry in the deleted setup.py)

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you run the unit tests and functional tests locally?
Did you add or update any necessary documentation?

copy-pr-bot · 2026-03-21T22:21:12Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

kajalj22 · 2026-03-21T22:21:25Z

/ok to test 85fbdad

kajalj22 · 2026-03-21T22:37:01Z

/ok to test e1b3b1b

kajalj22 · 2026-03-21T22:39:04Z

/ok to test 7741c47

kajalj22 · 2026-03-24T15:23:24Z

/ok to test a856567

kajalj22 · 2026-03-24T15:37:48Z

/ok to test 0c3dd8b

kajalj22 · 2026-03-25T17:08:34Z

/ok to test 142e33f

chtruong814 · 2026-03-28T12:12:36Z

@kajalj22 cc: @terrykong This PR helps get you further. The nv-grouped-gemm dependency needed to also be updated in the lock file.

Tests are still failing due to what is being installed. I'll see if I can help you get closer but a problem I think is that RL was installing Megatron-LM as a workspace, which will include it by default in the uv environment. It's also relying on some things that other projects typically do not rely on because they are technically not part of Megatron-Core itself like megatron.training.

The way RL is written or at least the way our tests are written, it assumes Megatron-LM is not an optional dependency. I may try to still make Megatron-LM a workspace still to keep things consistent but just point to code in MBridge instead. That may be easiest.

#2158

chtruong814 · 2026-03-29T13:53:33Z

I updated my PR to your PR to re-add Megatron-LM as a workspace. I am attempting to see if adding Megatron-LM as a symlink to MBridge's Megatron-LM reference will work better.

chtruong814 · 2026-03-30T00:56:42Z

Yeah the CI gets further with that change to still make Megatron-LM a workspace but symlinking the code from Megatron-Bridge's Megatron-LM submodule.

https://github.com/NVIDIA-NeMo/RL/actions/runs/23710543221/job/69073536800?pr=2158

However, it fails somewhere in Megatron-LM. I'd recommend check with Terry whether this approach seems feasible. If so, then resolving your merge conflicts and merging latest main branch of nemo-rl to your PR. There were some recent changes and perhaps those fixed this issue.

kajalj22 · 2026-03-30T15:57:41Z

Yeah the CI gets further with that change to still make Megatron-LM a workspace but symlinking the code from Megatron-Bridge's Megatron-LM submodule.

https://github.com/NVIDIA-NeMo/RL/actions/runs/23710543221/job/69073536800?pr=2158

However, it fails somewhere in Megatron-LM. I'd recommend check with Terry whether this approach seems feasible. If so, then resolving your merge conflicts and merging latest main branch of nemo-rl to your PR. There were some recent changes and perhaps those fixed this issue.

Sounds good, I will check with Terry, thank you for the help!

kajalj22 · 2026-04-06T16:09:01Z

/ok to test 6718424

kajalj22 · 2026-04-06T16:23:20Z

/ok to test 919b9ae

kajalj22 · 2026-05-15T15:54:04Z

/ok to test b465fdc

copy-pr-bot · 2026-05-15T20:18:39Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

yuki-97

thanks Kajal, some minor comments.

kajalj22 · 2026-05-18T15:46:58Z

/ok to test 3872c4e

Signed-off-by: Kajal Jain <kajalj@nvidia.com>

The MLM workspace setup.py previously provided a git-pinned source for nvidia-resiliency-ext. Without it, uv falls back to PyPI which has no cp313 wheels, breaking resolution. Signed-off-by: Kajal Jain <kajalj@nvidia.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

Signed-off-by: Kajal Jain <kajalj@nvidia.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

References to 3rdparty/Megatron-LM-workspace/Megatron-LM updated to 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM across runtime init, k8s configs, tests, docs, and tools. Signed-off-by: Kajal Jain <kajalj@nvidia.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

The removed Megatron-LM-workspace/setup.py directly listed onnxscript as a dependency. With that workspace gone, onnxscript is only available under megatron-core[dev] and megatron-core[lts] extras but the mcore extra installs plain megatron-core. transformer_engine imports onnxscript at module load time, causing --mcore-only tests to fail with ModuleNotFoundError. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

megatron-bridge unconditionally imports modelopt at module load time (auto_bridge.py imports is_quantized). The old MLM workspace's setup.py bundled nvidia-modelopt as a direct dependency; with that workspace removed, modelopt is no longer available under --extra mcore. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

The is_megatron_installed module was provided by the removed Megatron-LM-workspace package (via py_modules in setup.py). Replace the module imports with direct imports that fail fast under set -e. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

kajalj22 · 2026-05-18T21:57:54Z

/ok to test 38c16b6

github-actions Bot added the Documentation Improvements or additions to documentation label Mar 21, 2026

kajalj22 changed the title ~~Kajalj22/remove mlm workspace~~ fix: remove mlm workspace Mar 21, 2026

kajalj22 added CI:L2 Run doctests, unit tests, functional tests, and convergence tests CI:L1 Run doctests, unit tests, and functional tests and removed CI:L2 Run doctests, unit tests, functional tests, and convergence tests labels Mar 21, 2026

kajalj22 force-pushed the kajalj22/remove-mlm-workspace branch from 85fbdad to e1b3b1b Compare March 21, 2026 22:36

copy-pr-bot Bot had a problem deploying to nemo-ci March 21, 2026 22:37 Error

kajalj22 force-pushed the kajalj22/remove-mlm-workspace branch from e1b3b1b to 7741c47 Compare March 21, 2026 22:38

copy-pr-bot Bot had a problem deploying to nemo-ci March 21, 2026 22:40 Failure

copy-pr-bot Bot had a problem deploying to nemo-ci March 24, 2026 15:29 Failure

copy-pr-bot Bot had a problem deploying to nemo-ci March 24, 2026 15:38 Failure

copy-pr-bot Bot had a problem deploying to nemo-ci March 25, 2026 17:09 Failure

kajalj22 force-pushed the kajalj22/remove-mlm-workspace branch 3 times, most recently from 3ca7089 to 6718424 Compare April 6, 2026 16:07

copy-pr-bot Bot had a problem deploying to nemo-ci April 6, 2026 16:09 Failure

kajalj22 force-pushed the kajalj22/remove-mlm-workspace branch from 6718424 to 919b9ae Compare April 6, 2026 16:22

copy-pr-bot Bot temporarily deployed to public May 14, 2026 21:45 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci May 14, 2026 21:45 Error

copy-pr-bot Bot temporarily deployed to public May 14, 2026 21:45 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci May 14, 2026 22:18 Error

copy-pr-bot Bot temporarily deployed to nemo-ci May 14, 2026 22:40 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci May 14, 2026 23:43 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci May 15, 2026 01:16 Failure

copy-pr-bot Bot had a problem deploying to nemo-ci May 15, 2026 13:39 Failure

kajalj22 force-pushed the kajalj22/remove-mlm-workspace branch from eb45bc6 to b465fdc Compare May 15, 2026 15:53

copy-pr-bot Bot temporarily deployed to public May 15, 2026 15:54 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci May 15, 2026 15:54 Inactive

copy-pr-bot Bot temporarily deployed to public May 15, 2026 15:54 Inactive

copy-pr-bot Bot temporarily deployed to public May 15, 2026 15:59 Inactive

chtruong814 previously approved these changes May 15, 2026

View reviewed changes

yuki-97 reviewed May 16, 2026

View reviewed changes

Comment thread 3rdparty/Megatron-LM-workspace/is_megatron_installed.py Outdated

Comment thread pyproject.toml

Comment thread pyproject.toml

kajalj22 and others added 8 commits May 18, 2026 16:57

Remove megatron LM workspace

570b623

Signed-off-by: Kajal Jain <kajalj@nvidia.com>

Regenerate uv.lock after workspace removal

5bd7015

Signed-off-by: Kajal Jain <kajalj@nvidia.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

Remove stale CODEOWNERS entry for deleted MLM workspace

38c16b6

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Kajal Jain <kajalj@nvidia.com>

yuki-97 approved these changes May 19, 2026

View reviewed changes

terrykong mentioned this pull request Jul 16, 2026

build: install megatron-bridge from submodule; auto-flow its deps on bump #3235

Open

4 tasks

Uh oh!

Conversation

kajalj22 commented Mar 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Why

What changed

Before your PR is "Ready for review"

Uh oh!

copy-pr-bot Bot commented Mar 21, 2026

Uh oh!

kajalj22 commented Mar 21, 2026

Uh oh!

kajalj22 commented Mar 21, 2026

Uh oh!

kajalj22 commented Mar 21, 2026

Uh oh!

kajalj22 commented Mar 24, 2026

Uh oh!

kajalj22 commented Mar 24, 2026

Uh oh!

kajalj22 commented Mar 25, 2026

Uh oh!

chtruong814 commented Mar 28, 2026

Uh oh!

chtruong814 commented Mar 29, 2026

Uh oh!

chtruong814 commented Mar 30, 2026

Uh oh!

kajalj22 commented Mar 30, 2026

Uh oh!

kajalj22 commented Apr 6, 2026

Uh oh!

kajalj22 commented Apr 6, 2026

Uh oh!

kajalj22 commented May 15, 2026

Uh oh!

copy-pr-bot Bot commented May 15, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kajalj22 commented May 18, 2026

Uh oh!

kajalj22 commented May 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kajalj22 commented Mar 21, 2026 •

edited

Loading