Megatron-Bridge import example in launcher for Nemotron Super V3 [OMNIML-4827] - #1516
Merged
Conversation
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a bash entrypoint to run Megatron-Bridge HF→Megatron conversion, a Slurm job example invoking it, an optional ChangesMegatron-Bridge Checkpoint Import
🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tools/launcher/launch.py (1)
74-74: ⚡ Quick winAvoid hardcoding
relative_pathlength.This is easy to desync from
include_patternand can silently break packager mapping when entries change.Proposed refactor
-packager = run.PatternPackager( - include_pattern=[ +_include_patterns = [ "modules/Megatron-LM/megatron/*", "modules/Megatron-LM/examples/*", "modules/Megatron-LM/*.py", "modules/Megatron-Bridge/src/*", "modules/Megatron-Bridge/examples/*", "modules/Megatron-Bridge/pyproject.toml", "modules/Megatron-Bridge/README.md", "modules/Model-Optimizer/modelopt/*", "modules/Model-Optimizer/modelopt_recipes/*", "modules/Model-Optimizer/examples/*", "examples/*", "common/*", - ], - relative_path=[LAUNCHER_DIR] * 12, +] + +packager = run.PatternPackager( + include_pattern=_include_patterns, + relative_path=[LAUNCHER_DIR] * len(_include_patterns), )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tools/launcher/launch.py` at line 74, The hardcoded repeat count in the relative_path assignment (relative_path=[LAUNCHER_DIR] * 12) can get out of sync with include_pattern; change it to derive the repeat length from the include_pattern (e.g. relative_path=[LAUNCHER_DIR] * len(include_pattern) or equivalent) so that relative_path size always matches include_pattern length (reference variables: relative_path, LAUNCHER_DIR, include_pattern).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tools/launcher/common/megatron_bridge/import/import.sh`:
- Around line 45-49: The current import check only verifies importability and
can pick up an older installed megatron.bridge; change the logic so we force
local resolution by verifying the imported module's file path is inside
BRIDGE_DIR (or else install). Replace the simple import test with a Python check
that imports megatron.bridge, inspects megatron.bridge.__file__ (or uses
importlib to get the module path), resolves it to an absolute path and exits
nonzero if that path does not start with the absolute BRIDGE_DIR; keep the pip
install -e "${BRIDGE_DIR}" step to run when the path check fails. This ensures
the script uses the submodule under BRIDGE_DIR rather than an unrelated
site-packages version.
In
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_bridge_import.yaml`:
- Line 26: Replace the hardcoded OUTPUT_DIR value with a launcher interpolation
variable: remove "/scratchspace/megatron-bridge" and set OUTPUT_DIR to use the
shared launcher global var (e.g., <<global_vars.OUTPUT_DIR>>); update any
accompanying docs/examples to reference the global var if needed and ensure the
global_vars entry is defined in the launcher YAML set so the
megatron_bridge_import.yaml's OUTPUT_DIR uses <<global_vars.OUTPUT_DIR>> instead
of the literal path.
---
Nitpick comments:
In `@tools/launcher/launch.py`:
- Line 74: The hardcoded repeat count in the relative_path assignment
(relative_path=[LAUNCHER_DIR] * 12) can get out of sync with include_pattern;
change it to derive the repeat length from the include_pattern (e.g.
relative_path=[LAUNCHER_DIR] * len(include_pattern) or equivalent) so that
relative_path size always matches include_pattern length (reference variables:
relative_path, LAUNCHER_DIR, include_pattern).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3c1a8d8d-8cbf-4959-8dc3-ee644e8af2e1
📒 Files selected for processing (6)
.gitmodulestools/launcher/common/megatron_bridge/import/import.shtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_bridge_import.yamltools/launcher/launch.pytools/launcher/modules/Megatron-Bridgetools/launcher/modules/Megatron-LM
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1516 +/- ##
==========================================
+ Coverage 76.70% 76.77% +0.06%
==========================================
Files 477 477
Lines 51977 51977
==========================================
+ Hits 39869 39903 +34
+ Misses 12108 12074 -34
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tools/launcher/core.py (1)
139-145: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winDefine and maintain explicit module exports for this new public field.
GlobalVariablesgainedoutput_dirat Line 145, but the module still does not define an explicit__all__export surface. Please add/update__all__so this public API expansion is intentional and stable.As per coding guidelines, "Define the public API with
__all__at the top of each module and re-export viafrom .module import *in__init__.pyfiles."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tools/launcher/core.py` around lines 139 - 145, Add or update the module-level __all__ to explicitly declare the public API (including "GlobalVariables") so the newly added GlobalVariables.output_dir is intentionally exported; place __all__ at the top of the module and include all other public names from this file, and ensure the package __init__.py re-exports the module via "from .core import *" so consumers see the updated public surface.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tools/launcher/core.py`:
- Line 145: Add a regression test that verifies the new global variable
global_vars.output_dir is interpolated into both args and env like the existing
hf_model checks: create a test (e.g., test_global_vars_output_dir_interpolation)
that sets global_vars.output_dir to a sample value, constructs a launcher/job
spec with args and env entries containing the token <<global_vars.output_dir>>,
runs the same interpolation routine used by the hf_model tests (the helper or
method those tests call), and assert that the token is replaced with the sample
path in both the args list and environment dict; mirror the structure and
assertions of the hf_model interpolation tests so the new test validates both
args and env interpolation for output_dir.
---
Outside diff comments:
In `@tools/launcher/core.py`:
- Around line 139-145: Add or update the module-level __all__ to explicitly
declare the public API (including "GlobalVariables") so the newly added
GlobalVariables.output_dir is intentionally exported; place __all__ at the top
of the module and include all other public names from this file, and ensure the
package __init__.py re-exports the module via "from .core import *" so consumers
see the updated public surface.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e25f9ee3-5962-4d2a-a763-49b32fcd1609
📒 Files selected for processing (2)
tools/launcher/core.pytools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_bridge_import.yaml
🚧 Files skipped from review as they are similar to previous changes (1)
- tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_bridge_import.yaml
| hf_model: str = None | ||
| hf_data: str = None | ||
| hf_local: str = None | ||
| output_dir: str = None |
Contributor
There was a problem hiding this comment.
Add a regression test for <<global_vars.output_dir>> interpolation.
Line 145 introduces a new interpolated global var, but there’s no corresponding test case in the provided tools/launcher/tests/test_core.py snippets. Please add one test mirroring the existing hf_model interpolation checks for args/environment using output_dir.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tools/launcher/core.py` at line 145, Add a regression test that verifies the
new global variable global_vars.output_dir is interpolated into both args and
env like the existing hf_model checks: create a test (e.g.,
test_global_vars_output_dir_interpolation) that sets global_vars.output_dir to a
sample value, constructs a launcher/job spec with args and env entries
containing the token <<global_vars.output_dir>>, runs the same interpolation
routine used by the hf_model tests (the helper or method those tests call), and
assert that the token is replaced with the sample path in both the args list and
environment dict; mirror the structure and assertions of the hf_model
interpolation tests so the new test validates both args and env interpolation
for output_dir.
jenchen13
enabled auto-merge (squash)
May 27, 2026 15:12
jenchen13
changed the title
Contributor
|
jenchen13
added a commit
that referenced
this pull request
May 27, 2026
### What does this PR do? Type of change: New feature Add Megatron-Bridge import example in launcher for Nemotron Super V3 ### Usage ```python # Usage: # update .env-slurm with environment variables # cd tools/launcher # uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/megatron_bridge_import.yaml --yes ``` ### Testing <!-- Mention how have you tested your change if applicable. --> ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A <!--- Mandatory --> - Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> - Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run `/claude review`. NVIDIA org members can self-trigger for complex changes; orthogonal to CodeRabbit. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added Megatron-Bridge Hugging Face→Megatron checkpoint import (CPU-capable entrypoint) and an example pipeline for importing the Nvidia Nemotron-3 Super 120B model. * Exposed an `output_dir` global variable for pipeline interpolation. * **Chores** * Updated pinned Megatron-LM module. * Extended .gitignore to cover `.env*` variants. <!-- review_stack_entry_start --> [](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1516?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1516?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack){kind=link}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: New feature
Add Megatron-Bridge import example in launcher for Nemotron Super V3
Usage
Testing
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information
Summary by CodeRabbit
New Features
Chores