Add HybridModel MBridge support for nemo:26.08#2005
Conversation
59105eb to
a3c2ecf
Compare
|
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:
📝 WalkthroughWalkthroughMegatron integration adds HybridModel support, compatibility with newer Mamba convolution layouts, renamed hybrid layer-pattern configuration, updated statistics and tests, revised packed-dataset sampling, and related example, warning, pruning-log, and container documentation changes. ChangesMegatron hybrid model integration
Runtime and data behavior updates
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Loader as load_mbridge_model_from_hf
participant Provider as HybridModelProvider
participant Stack as get_te_hybrid_stack_spec
participant Model as HybridModel
Loader->>Provider: detect hybrid provider
Loader->>Stack: configure hybrid stack
Stack-->>Loader: return ModuleSpec
Loader->>Model: validate loaded model
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/claude review |
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Full review. 7 files changed (+53/-21). Reviewed all modelopt/ and examples/ source files; CHANGELOG.rst and tools/launcher/.../import.sh are doc/version-bump only.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 2
This is a clean, focused change adding HybridModel / HybridModelProvider support for nemo:26.08 (where Nemotron-H initializes as a plain HybridModel rather than the MambaModel subclass).
Strengths:
- Plugin laziness preserved:
HybridModel/HybridModelProviderimports are gated behindtry/exceptwith anHAS_HYBRIDflag in bothnas/plugins/megatron.pyandutils/plugins/mbridge.py; users without nemo:26.08 extras are unaffected. get_te_hybrid_stack_specmirrors the existingget_te_mamba_stack_specpattern (deepcopy + SequentialMLP replacement to avoid the grouped-GEMM dependency for pruning), and is correctly added to__all__.- Type hints, the
isinstanceassertion, and the return-tuple docstring are all updated consistently to include the newHybridModeltypes. - The
sample_multiplier8→16 bump indataset_utils.pycorrectly threads the value into the warning message so it no longer hardcodes8x. distill.pymoving eval args intoValidationConfigmatches the code comment's stated migration path.
Non-blocking suggestions (inline):
quantize.py:--quant_cfgdefault"fp8"→None— with neither--recipenor--quant_cfgset, the failure surfaces as a misleading "Unsupported --quant_cfg 'None'" deep inget_quant_config; suggest an explicitparser.errorguard.mbridge.py: verifyMambaModelProvideris not a subclass ofHybridModelProvider— if it is, theisinstance(HybridModelProvider)branch would shadow the Mamba path (branch ordering).
Risk: low. Additive, well-gated, and consistent with existing patterns. Author notes it was tested with Nemotron-3-Nano PTQ on MBridge main, with GH CICD tests to follow at the nemo:26.08 release.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2005 +/- ##
==========================================
+ Coverage 75.87% 77.81% +1.93%
==========================================
Files 518 518
Lines 58544 58574 +30
==========================================
+ Hits 44422 45578 +1156
+ Misses 14122 12996 -1126
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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: 2
🧹 Nitpick comments (2)
modelopt/torch/utils/plugins/mbridge.py (1)
38-46: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winClarify the container baseline for HybridModel support.
The hybrid provider path is annotated
nemo:26.08+, while the release note and import helper advertisenemo:26.04+without distinguishing generic Bridge support from HybridModel support. Document 26.08 as the hybrid requirement, or explicitly state that 26.04 applies only to non-hybrid workflows.
modelopt/torch/utils/plugins/mbridge.py#L38-L46: document the hybrid-specific 26.08+ requirement.CHANGELOG.rst#L19-L19: distinguish the generic 26.04 baseline from HybridModel support.tools/launcher/common/megatron_bridge/import/import.sh#L19-L19: clarify whether this script supports hybrid models on 26.04.🤖 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 `@modelopt/torch/utils/plugins/mbridge.py` around lines 38 - 46, Clarify the version baseline across all affected documentation: in modelopt/torch/utils/plugins/mbridge.py lines 38-46, document that HybridModelProvider and HybridModel require nemo:26.08+; in CHANGELOG.rst line 19, distinguish the generic Megatron Bridge 26.04+ baseline from hybrid support; and in tools/launcher/common/megatron_bridge/import/import.sh line 19, explicitly state whether the 26.04+ script supports hybrid models.examples/megatron_bridge/quantize.py (1)
127-127: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winValidate the required quantization source at the CLI boundary.
With this new default, omitting both
--recipeand--quant_cfgreaches the unsupported-config error inget_quant_config()instead of producing a clear parser error. If removing the implicit FP8 preset is intentional, require one of these options during argument parsing; otherwise add an explicit fallback.Proposed validation
args = parser.parse_args() + if args.recipe is None and args.quant_cfg is None: + parser.error("one of --recipe or --quant_cfg is required") + if args.moe_calib_experts_ratio is not None:🤖 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 `@examples/megatron_bridge/quantize.py` at line 127, Update the CLI argument definition in the quantization entrypoint so omitting both --recipe and --quant_cfg is rejected during argument parsing when no implicit FP8 preset exists. Require one of the two options together while preserving their existing individual handling, or restore an explicit fallback configuration if that preset is still intended.
🤖 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 `@modelopt/torch/utils/dataset_utils.py`:
- Around line 840-842: Update the packed-mode oversampling documentation in
get_dataset_dataloader and the related packing test comment from 4x and 8x to
16x, respectively, so all references match the sample_multiplier behavior.
- Around line 865-867: Update the pack=True exhaustion message near the
row-count diagnostic to prioritize shortening max_sample_length or supplying
longer, more token-dense samples; remove the broad recommendation to increase
num_samples, or qualify it only for sources that can provide additional useful
content.
---
Nitpick comments:
In `@examples/megatron_bridge/quantize.py`:
- Line 127: Update the CLI argument definition in the quantization entrypoint so
omitting both --recipe and --quant_cfg is rejected during argument parsing when
no implicit FP8 preset exists. Require one of the two options together while
preserving their existing individual handling, or restore an explicit fallback
configuration if that preset is still intended.
In `@modelopt/torch/utils/plugins/mbridge.py`:
- Around line 38-46: Clarify the version baseline across all affected
documentation: in modelopt/torch/utils/plugins/mbridge.py lines 38-46, document
that HybridModelProvider and HybridModel require nemo:26.08+; in CHANGELOG.rst
line 19, distinguish the generic Megatron Bridge 26.04+ baseline from hybrid
support; and in tools/launcher/common/megatron_bridge/import/import.sh line 19,
explicitly state whether the 26.04+ script supports hybrid models.
🪄 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: bb158952-0cce-421c-a2e3-c8e2648f838c
📒 Files selected for processing (7)
CHANGELOG.rstexamples/megatron_bridge/distill.pyexamples/megatron_bridge/quantize.pymodelopt/torch/nas/plugins/megatron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/plugins/mbridge.pytools/launcher/common/megatron_bridge/import/import.sh
a3c2ecf to
09cadd3
Compare
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: 2
🤖 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 `@modelopt/torch/utils/plugins/mbridge.py`:
- Around line 102-107: Update the provider handling around HybridModelProvider
and MambaModelProvider so get_te_hybrid_stack_spec and get_te_mamba_stack_spec
do not overwrite specifications already supplied through provider_overrides.
Preserve each provider’s existing hybrid_stack_spec or mamba_stack_spec when
present; only install the generic spec when the corresponding value is unset, or
merge MoE-related changes while retaining the provider-specific layer
composition and explicit override precedence.
- Around line 62-64: Update the fourth return-type element in the relevant
mbridge function to include the full VLM wrapper type, using the existing
MegatronModule base type or an explicit supported wrapper union alongside
GPTModel, MambaModel, and HybridModel. Keep the implementation’s VLM return
behavior unchanged.
🪄 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: 1f3c3691-4eee-421c-90d9-5245861be04e
📒 Files selected for processing (8)
CHANGELOG.rstexamples/megatron_bridge/distill.pyexamples/megatron_bridge/quantize.pymodelopt/torch/nas/plugins/megatron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/plugins/mbridge.pytests/unit/torch/utils/test_dataset_utils.pytools/launcher/common/megatron_bridge/import/import.sh
🚧 Files skipped from review as they are similar to previous changes (6)
- tools/launcher/common/megatron_bridge/import/import.sh
- examples/megatron_bridge/quantize.py
- modelopt/torch/utils/dataset_utils.py
- examples/megatron_bridge/distill.py
- CHANGELOG.rst
- modelopt/torch/nas/plugins/megatron.py
09cadd3 to
9019122
Compare
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
9019122 to
a1e035a
Compare
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)
examples/megatron_bridge/prune_minitron.py (1)
553-559: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRead hybrid metadata from the pruned language model.
For VLMs,
unwrapped_modelis the outer wrapper whilelanguage_modelis the pruned HybridModel. This new provider path readshybrid_layer_patternfromunwrapped_modelat Lines 556 and 652, which can raiseAttributeErroror write stale metadata. Uselanguage_modelconsistently.Proposed fix
- if hasattr(unwrapped_model, "hybrid_override_pattern") + if hasattr(language_model, "hybrid_override_pattern") else "hybrid_layer_pattern" ) - setattr(provider, hybrid_key, getattr(unwrapped_model, hybrid_key)) + setattr(provider, hybrid_key, getattr(language_model, hybrid_key)) ... - hf_cfg.hybrid_override_pattern = getattr(unwrapped_model, hybrid_key) + hf_cfg.hybrid_override_pattern = getattr(language_model, hybrid_key)Also applies to: 649-651
🤖 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 `@examples/megatron_bridge/prune_minitron.py` around lines 553 - 559, Update the hybrid provider metadata assignments in the paths around the provider handling at lines 553 and 649 to read the selected hybrid key from language_model instead of unwrapped_model. Keep the existing hybrid_override_pattern versus hybrid_layer_pattern key selection and assign the value from the pruned HybridModel consistently.
🤖 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 `@tests/_test_utils/torch/megatron/models.py`:
- Around line 442-445: Update the validation near hybrid_layer_pattern
construction to explicitly reject any "|" characters before checking its length.
Remove the pipe-stripping validation and assert that the pattern is pipe-free,
while preserving the existing num_layers length check for valid patterns.
---
Outside diff comments:
In `@examples/megatron_bridge/prune_minitron.py`:
- Around line 553-559: Update the hybrid provider metadata assignments in the
paths around the provider handling at lines 553 and 649 to read the selected
hybrid key from language_model instead of unwrapped_model. Keep the existing
hybrid_override_pattern versus hybrid_layer_pattern key selection and assign the
value from the pruned HybridModel consistently.
🪄 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: 02808af7-98f7-41a9-94fd-060aa447679a
📒 Files selected for processing (19)
CHANGELOG.rstexamples/megatron_bridge/distill.pyexamples/megatron_bridge/prune_minitron.pyexamples/megatron_bridge/quantize.pymodelopt/torch/nas/plugins/megatron.pymodelopt/torch/nas/plugins/megatron_model_stats.pymodelopt/torch/prune/plugins/mcore_minitron.pymodelopt/torch/utils/dataset_utils.pymodelopt/torch/utils/import_utils.pymodelopt/torch/utils/plugins/mbridge.pytests/_test_utils/torch/megatron/models.pytests/_test_utils/torch/megatron/utils.pytests/gpu_megatron/torch/export/test_megatron_importer.pytests/gpu_megatron/torch/nas/plugins/test_megatron_mamba_dynamic_modules.pytests/gpu_megatron/torch/nas/plugins/test_megatron_model_stats.pytests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.pytests/gpu_megatron/torch/quantization/plugins/test_megatron.pytests/unit/torch/utils/test_dataset_utils.pytools/launcher/common/megatron_bridge/import/import.sh
🚧 Files skipped from review as they are similar to previous changes (7)
- CHANGELOG.rst
- tools/launcher/common/megatron_bridge/import/import.sh
- examples/megatron_bridge/quantize.py
- modelopt/torch/utils/dataset_utils.py
- tests/unit/torch/utils/test_dataset_utils.py
- examples/megatron_bridge/distill.py
- modelopt/torch/utils/plugins/mbridge.py
a1e035a to
e348e26
Compare
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
e348e26 to
ed67baf
Compare
1 participant
Description
MBridge main (nemo:26.08) will initialize Nemotron-H as HybridModel instead of MambaModel (subclassed of HybridModel). Also make minimum nemo container 26.04
Testing
Tested Nemotron-3-Nano PTQ with MBridge main (fails otherwise)
Tested locally
tests/gpu_megatronandtests/examples/megatron_bridgewithnemo:26.06.01+ Mount latest MBridge/McoreGH CICD tests will be added with nemo:26.08 release
Summary by CodeRabbit
--quant_cfgnow defaults toNoneunless explicitly set (or via--recipe).nemo:26.04or newer (nemo:26.06recommended).