MBridge VLM distillation / QAD support#1938
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
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 VLM-aware Megatron-Bridge distillation, a distilled-checkpoint Hugging Face exporter, conditional checkpoint state restoration, updated quantized export workflows, hybrid-attention pruning support, and expanded tests and documentation. ChangesMegatron-Bridge distillation and export flow
Hybrid-attention pruning support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant distill.py
participant run_provider
participant convert_to_distillation_provider
participant save_vlm_to_hf
participant export_llm_to_hf
distill.py->>run_provider: Load student checkpoint weights
run_provider->>convert_to_distillation_provider: Run deferred KD conversion
convert_to_distillation_provider-->>distill.py: Return distilled model
distill.py->>save_vlm_to_hf: Export distilled VLM
distill.py->>export_llm_to_hf: Export distilled LLM
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
kevalmorabia97
changed the title
kevalmorabia97
requested review from
AAnoosheh,
ChenhanYu,
danielkorzekwa,
j-rausch and
jenchen13
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1938 +/- ##
==========================================
- Coverage 77.92% 73.81% -4.11%
==========================================
Files 519 522 +3
Lines 58217 59487 +1270
==========================================
- Hits 45366 43913 -1453
- Misses 12851 15574 +2723
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:
|
|
/claude review |
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/examples/megatron_bridge/test_quantize_export.py (1)
15-83: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winKeep the VLM export coverage
test_qad.pyonly reachesexport_quantized_megatron_to_hf.pyon the non-VLM path; the VLM branch returns before export, so it does not replace the removed case.🤖 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 `@tests/examples/megatron_bridge/test_quantize_export.py` around lines 15 - 83, The VLM export path still needs direct coverage because the QAD flow exits before reaching export_quantized_megatron_to_hf.py, so this test should keep validating the quantize-and-export bridge for the VLM-specific Megatron checkpoint. Preserve the existing quantize/export steps in test_quantize_and_export, using the same helpers like create_tiny_qwen3_dir, extend_cmd_parts, and run_example_command, so the removed coverage is not lost.
🧹 Nitpick comments (1)
tests/examples/megatron_bridge/test_prune_minitron.py (1)
26-40: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate helper across test files.
_qwen35_moe_bridge_supported()is duplicated verbatim intests/examples/megatron_bridge/test_qad.py. Consider moving it to a shared location (e.g. a localconftest.pyintests/examples/megatron_bridge/) to avoid drift between the two copies.♻️ Suggested consolidation
# tests/examples/megatron_bridge/conftest.py from megatron.bridge.models.conversion import model_bridge def qwen35_moe_bridge_supported() -> bool: """Whether MBridge supports Qwen3.5-MoE per-expert weight assembly, i.e. nemo:26.08+.""" try: return hasattr(model_bridge, "_fuse_per_expert_hf_weight") except Exception: return FalseThen import
qwen35_moe_bridge_supportedfrom the shared module in both test files.🤖 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 `@tests/examples/megatron_bridge/test_prune_minitron.py` around lines 26 - 40, The helper _qwen35_moe_bridge_supported() is duplicated in multiple Megatron bridge test files, so consolidate it into a shared test fixture/module such as a local conftest.py under tests/examples/megatron_bridge/. Move the logic that checks model_bridge for _fuse_per_expert_hf_weight into a single shared function (for example qwen35_moe_bridge_supported), then import and use that shared helper from both test_prune_minitron.py and test_qad.py to keep the behavior in sync.
🤖 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 `@examples/megatron_bridge/distill.py`:
- Around line 377-391: In the VLM export path inside distill.py, the call to
save_vlm_to_hf is using the wrong source HF identifier for config/tokenizer
lookup. Update the export block under if args.hf_export_path and is_vlm so it
passes args.student_hf_path instead of args.student_hf_model, matching the
standalone export_distilled_megatron_to_hf flow and the intended VLM source path
semantics.
In `@examples/megatron_bridge/quantize.py`:
- Around line 417-420: The post-save message in quantize.py is too generic
because it tells all users to convert with export_quantized_megatron_to_hf.py
even though VLM checkpoints remain Megatron-only. Update the print_rank_0 hint
near the quantized model save path to branch on is_vlm (or explicitly mention
the VLM exception) so non-VLM users see the conversion guidance while VLM users
are warned that the conversion path is unsupported.
---
Outside diff comments:
In `@tests/examples/megatron_bridge/test_quantize_export.py`:
- Around line 15-83: The VLM export path still needs direct coverage because the
QAD flow exits before reaching export_quantized_megatron_to_hf.py, so this test
should keep validating the quantize-and-export bridge for the VLM-specific
Megatron checkpoint. Preserve the existing quantize/export steps in
test_quantize_and_export, using the same helpers like create_tiny_qwen3_dir,
extend_cmd_parts, and run_example_command, so the removed coverage is not lost.
---
Nitpick comments:
In `@tests/examples/megatron_bridge/test_prune_minitron.py`:
- Around line 26-40: The helper _qwen35_moe_bridge_supported() is duplicated in
multiple Megatron bridge test files, so consolidate it into a shared test
fixture/module such as a local conftest.py under
tests/examples/megatron_bridge/. Move the logic that checks model_bridge for
_fuse_per_expert_hf_weight into a single shared function (for example
qwen35_moe_bridge_supported), then import and use that shared helper from both
test_prune_minitron.py and test_qad.py to keep the behavior in sync.
🪄 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: a418a2b7-0f1b-4330-a397-723d252f5670
📒 Files selected for processing (14)
CHANGELOG.rstexamples/megatron_bridge/README.mdexamples/megatron_bridge/_distillation_provider.pyexamples/megatron_bridge/distill.pyexamples/megatron_bridge/export_distilled_megatron_to_hf.pyexamples/megatron_bridge/export_quantized_megatron_to_hf.pyexamples/megatron_bridge/quantize.pyexamples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.mdexamples/megatron_bridge/tutorials/README.mdmodelopt/torch/utils/plugins/mbridge.pytests/examples/megatron_bridge/test_distill.pytests/examples/megatron_bridge/test_prune_minitron.pytests/examples/megatron_bridge/test_qad.pytests/examples/megatron_bridge/test_quantize_export.py
![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: Reviewed all 14 changed files, prioritizing new/modified source (_distillation_provider.py, distill.py, export_distilled_megatron_to_hf.py, quantize.py, modelopt/torch/utils/plugins/mbridge.py), then tests and README/CHANGELOG. Examples + tests PR plus one backward-compatible signature change (load_modelopt_megatron_checkpoint gains restore_modelopt_state: bool = True; default preserves prior behavior).
Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 1 (posted inline) — the inline VLM export branch in distill.py (--hf_export_path + VLM) is not exercised by CI; test_distill_vlm validates only the standalone export_distilled_megatron_to_hf.py path.
Verified: New public imports resolve (copy_hf_ckpt_remote_code, warn_rank_0/print_args). Removed _patch_qwen35_moe_sequential_expert_mappings and the deleted export.py name have no remaining references in examples/ or docs/. The load_modelopt_megatron_checkpoint change is additive and the new restore_modelopt_state=False call site is correct. dist.cleanup() is guarded so the double-cleanup is safe. The teacher-mirroring __setattr__ foot-gun in _distillation_provider.py is handled via object.__setattr__.
Already flagged by CodeRabbit (not duplicated; worth confirming before merge): inline VLM export passes student_hf_model instead of student_hf_path to save_vlm_to_hf; quantize.py post-save message points all users at the export script though VLM quantized export is unsupported; duplicated _qwen35_moe_bridge_supported() helper across two test files.
Risk: Low. Example/test-scoped with documented temporary WARs (nemo:26.08 container) and one backward-compatible plugin signature change. No mode-registration, config-schema, or checkpoint-state-dict changes affecting existing optimized checkpoints.
|
Addressed review comments in 623b6c9:
All example-test changes collect cleanly; the code fixes are message/arg-only (no behavior change for the default |
🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
gh pr view 1938 --repo NVIDIA/Model-Optimizer --json commits --jq '.commits[] | select(.oid | startswith("623b6c9")) | .oid'
gh api repos/NVIDIA/Model-Optimizer/commits/623b6c9 --jq '.files[].filename' 2>/dev/null || trueLength of output: 647
Thanks also for confirming the ✅ Action performedReview finished.
|
kevalmorabia97
force-pushed
the
kmorabia/vlm-distill
branch
from
623b6c9 to
3b7106a
Compare
- distill.py: distill only the language_model submodule of a VLM (vision tower / projector untouched); reuses the LLM training path for everything else. - Add export_distilled_megatron_to_hf.py to convert a distilled Megatron checkpoint (any iteration) to HF -- LLM/Puzzletron via export_ckpt, VLM via in-memory reassembly (vision from HF + distilled LM from the checkpoint). - Rename export.py -> export_quantized_megatron_to_hf.py; distill.py's final export reuses the shared export helpers. - mbridge: load_modelopt_megatron_checkpoint gains restore_modelopt_state (weights-only reload of a full-precision distilled student, no teacher). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…helper - distill.py: VLM HF export uses --student_hf_path (the VLM config/tokenizer source), not --student_hf_model (which is unused for VLMs and only coincides via its default). - quantize.py: only point non-VLM checkpoints at export_quantized_megatron_to_hf.py; quantized-VLM HF export is unsupported, so say so for VLMs. - Move the duplicated _qwen35_moe_bridge_supported() helper into tests/_test_utils/examples/megatron_bridge.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
force-pushed
the
kmorabia/vlm-distill
branch
from
3b7106a to
1fee603
Compare
|
Tested on Cosmos-Reason2-2B:
Note: KD/QAD only applied only for short run of 50 iters Results: Key takeaways: particularly for NVFP4, QAD recovers drop in performance. FP8 less impacted overall, difference to BF16 is noise. |
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
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
🤖 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/prune/plugins/mcore_minitron.py`:
- Around line 423-454: Update _temporarily_pruned so the original all_layers
list and layer numbers are restored onto model.decoder.layers before calling
sample(model, sample_func=max). Then sample the reattached full layer list to
restore max-subnet state, followed by restoring per-layer patterns and hybrid
configuration as currently handled.
🪄 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: 5cfd7fbc-1d97-4c0e-93b6-c1ff5b3e4b8c
📒 Files selected for processing (6)
examples/megatron_bridge/README.mdexamples/megatron_bridge/prune_minitron.pymodelopt/torch/nas/plugins/megatron_model_stats.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/gpu_megatron/torch/nas/plugins/test_megatron_model_stats.pytests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/megatron_bridge/README.md
| @contextmanager | ||
| def _temporarily_pruned(self, export_config: dict): | ||
| """Prune to ``export_config`` for the duration of the block, then restore the max subnet. | ||
|
|
||
| Used to score a candidate subnet without permanently mutating the model. ``_prune`` drops | ||
| layers and slices per-layer patterns (config lists + hybrid string) in place; ``sample(max)`` | ||
| only restores hparam-controlled state, so those must be snapshotted and restored here — else | ||
| the mutations compound across candidates and corrupt the final export. | ||
| """ | ||
| model = self.model | ||
| all_layers = model.decoder.layers | ||
| start_layer_number = all_layers[0].layer_number | ||
| hybrid_key = _get_hybrid_pattern_key(model) | ||
| saved_patterns = { | ||
| attr: getattr(model.config, attr) | ||
| for attr in _PER_LAYER_CONFIG_PATTERNS | ||
| if isinstance(getattr(model.config, attr, None), (list, tuple)) | ||
| } | ||
| saved_hybrid = getattr(model, hybrid_key) if hybrid_key else None | ||
| try: | ||
| self._prune(export_config) | ||
| yield | ||
| finally: | ||
| # reset to max subnet, revert dropped layers, and restore in-place-sliced patterns | ||
| sample(model, sample_func=max) | ||
| for offset, layer in enumerate(all_layers): | ||
| layer.layer_number = start_layer_number + offset | ||
| model.decoder.layers = all_layers | ||
| for attr, pattern in saved_patterns.items(): | ||
| setattr(model.config, attr, pattern) | ||
| if hybrid_key: | ||
| setattr(model, hybrid_key, saved_hybrid) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for tied/shared hparam mechanisms for per-layer width dims in Megatron NAS plugins
rg -nP -C4 'class .*Hparam' modelopt/torch/nas/plugins/megatron.py modelopt/torch/opt/dynamic.py
rg -nP -C3 '\btied\b|shared.*hparam|_hp_registry' modelopt/torch/nas/plugins/megatron.py modelopt/torch/opt/dynamic.pyRepository: NVIDIA/Model-Optimizer
Length of output: 1826
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== megatron.py around shared hparam setup =="
sed -n '1210,1285p' modelopt/torch/nas/plugins/megatron.py
echo
echo "== dynamic.py hparam registration/shared semantics =="
rg -n -C4 'shared hidden_size|set the shared hidden_size hparam|_register_hparam|def get_hparam|def _register_hparam|share' modelopt/torch/opt/dynamic.py modelopt/torch/nas/plugins/megatron.py
echo
echo "== dynamic.py hparam-related class/method definitions =="
ast-grep outline modelopt/torch/opt/dynamic.py --view expanded | sed -n '1,240p'Repository: NVIDIA/Model-Optimizer
Length of output: 34413
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== search for sample(max) behavior and layer traversal =="
rg -n -C4 'def sample\(|sample_func=max|named_modules\(|decoder.layers' modelopt/torch/opt/dynamic.py modelopt/torch/prune/plugins/mcore_minitron.py
echo
echo "== inspect _temporarily_pruned surrounding code =="
sed -n '380,470p' modelopt/torch/prune/plugins/mcore_minitron.pyRepository: NVIDIA/Model-Optimizer
Length of output: 16457
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== dynamic.py sample/export semantics =="
rg -n -C6 'def sample\(|sample_func|max|def export\(' modelopt/torch/opt/dynamic.py
echo
echo "== dynamic.py around sample implementation =="
sed -n '320,420p' modelopt/torch/opt/dynamic.py
echo
echo "== dynamic.py around export implementation =="
sed -n '510,580p' modelopt/torch/opt/dynamic.pyRepository: NVIDIA/Model-Optimizer
Length of output: 11040
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== inspect hparam classes in dynamic.py =="
rg -n -C4 'class .*Hparam|class TracedHp|class .*Hp' modelopt/torch/opt/dynamic.py modelopt/torch/nas/plugins/megatron.py
echo
echo "== inspect _register_hparam and named_hparams traversal =="
sed -n '378,420p' modelopt/torch/opt/dynamic.py
sed -n '1160,1205p' modelopt/torch/opt/dynamic.pyRepository: NVIDIA/Model-Optimizer
Length of output: 7759
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read only the sample helper in dynamic.py if present, and the Hparam classes.
rg -n -C8 'def sample\(|sample_func=max|class TracedHp|class Hparam|def named_hparams\(' modelopt/torch/opt/dynamic.pyRepository: NVIDIA/Model-Optimizer
Length of output: 1661
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== sample helper in dynamic.py =="
rg -n -C12 'def sample\(|sample_func=max' modelopt/torch/opt/dynamic.py
echo
echo "== Hparam classes in dynamic.py =="
rg -n -C8 'class Hparam|class TracedHp|class .*Hp' modelopt/torch/opt/dynamic.pyRepository: NVIDIA/Model-Optimizer
Length of output: 194
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate sample() definitions and callers =="
rg -n -C4 'def sample\(' modelopt
echo
echo "== locate sample(model, sample_func=max) imports =="
rg -n -C4 'sample\(model, sample_func=max\)|from .* import sample|import .*sample' modelopt/torch/prune/plugins/mcore_minitron.py modeloptRepository: NVIDIA/Model-Optimizer
Length of output: 10977
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== inspect top of mcore_minitron.py for sample import =="
sed -n '1,120p' modelopt/torch/prune/plugins/mcore_minitron.pyRepository: NVIDIA/Model-Optimizer
Length of output: 4372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== inspect modelopt/torch/nas/utils.py sample implementation =="
sed -n '1,240p' modelopt/torch/nas/utils.py
echo
echo "== inspect selection semantics in modelopt/torch/nas/search_space.py =="
sed -n '1,180p' modelopt/torch/nas/search_space.pyRepository: NVIDIA/Model-Optimizer
Length of output: 16576
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== inspect _DynamicMLP and _DynamicMoELayer setup =="
sed -n '316,390p' modelopt/torch/nas/plugins/megatron.py
sed -n '730,855p' modelopt/torch/nas/plugins/megatron.pyRepository: NVIDIA/Model-Optimizer
Length of output: 9762
Restore the full layer list before sample(max)
sample(model, sample_func=max) only walks reachable modules, so the detached layers in all_layers keep their pruned per-layer widths and get reattached in a stale state. Swap the restore order so model.decoder.layers is set back first, then sample the max subnet.
🤖 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/prune/plugins/mcore_minitron.py` around lines 423 - 454,
Update _temporarily_pruned so the original all_layers list and layer numbers are
restored onto model.decoder.layers before calling sample(model,
sample_func=max). Then sample the reattached full layer list to restore
max-subnet state, followed by restoring per-layer patterns and hybrid
configuration as currently handled.
3 participants
What
Adds VLM (e.g. Qwen3.5-VL, Gemma3-VL) knowledge-distillation / QAD support to the Megatron-Bridge examples:
distill.pydistills only the language model submodule (vision tower + projector untouched), reusing the LLM training path.export_distilled_megatron_to_hf.pyconverts a distilled Megatron checkpoint (any iteration) to HF. Required especially for VLM distilled ckpt as it only has LM weights so we need to initialize full VLM, swap LLM weights then save to HFexport.py→export_quantized_megatron_to_hf.py.Related upstream Megatron-Bridge PRs to be available in nemo:26.08 container:
DistillationProvidersubmodule distillation (non-blocking; added temporary WAR)Testing
nemo:26.06(distill → separate HF export; LLM + VLM, incl. TP→TP/PP reshard).test_distill_vlmruns the export script as a CI e2e step.Summary by CodeRabbit
Summary
New Features
Bug Fixes
Documentation
Tests