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 unit tests (beta)
Comment |
kevalmorabia97
changed the title
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1938 +/- ##
==========================================
- Coverage 77.87% 75.80% -2.08%
==========================================
Files 522 525 +3
Lines 58452 61240 +2788
==========================================
+ Hits 45522 46423 +901
- Misses 12930 14817 +1887
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.
|
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. |
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
- 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>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
3ea6f6b to
cbddfbe
Compare
danielkorzekwa
left a comment
danielkorzekwa
left a comment
There was a problem hiding this comment.
Good work. I reviewed the READE first, let's resolves comments first, then I can review the code as well
| `--student_hf_model` is also needed for heterogeneous (Puzzletron/NAS) students whose architecture differs from `--student_hf_path`; for standard (homogeneous) models it defaults to `--student_hf_path`, and it is unused for VLMs. | ||
|
|
||
| **Separate conversion** -- convert any saved iteration using the Megatron-Bridge conversion script: | ||
| **Separate conversion** -- convert **any** saved iteration (intermediate or final) with [export_distilled_megatron_to_hf.py](export_distilled_megatron_to_hf.py): |
There was a problem hiding this comment.
why to emphasize it?
There was a problem hiding this comment.
The bold any is intentional — it contrasts with the Inline path just above, which only exports the final checkpoint. It flags that the separate script can export any intermediate iteration too. Happy to drop the emphasis if you'd prefer.
|
|
||
| ## Resources | ||
|
|
||
| - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699) |
There was a problem hiding this comment.
why to have roadmap here under resources, we do not have such resources section at other example readme files? we point road to a specific issue/1699 - will the roadmap always stay under this issue?
There was a problem hiding this comment.
The ## Resources + 📅 Roadmap footer is a shared convention across all example READMEs (hf_ptq, pruning, llm_distill, llm_qat, diffusers, onnx_ptq, torch_onnx, torch_trt, speculative_decoding), and #1699 is the canonical Model-Optimizer roadmap issue — so I'd keep it here for consistency with the other examples rather than diverge.
Restructure the Distillation section: non-quantized methods (real-data / VLM) followed by their HF conversion, then QAD (quantized, exported via export_quantized) and Slurm; move Sanity-Check Generation last; add a workflows intro; rename MEGATRON_BRIDGE_SRC_DIR; and clarify the --student_hf_model, --trust_remote_code, and export --pp_size notes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
Thanks for the review! Addressed all README comments in d5e97d8 (replied inline; resolved the addressed threads). Two left open for your call:
Ready for the code review whenever you are. |
…mporarily_pruned _prune sets width hparams on all layers before dropping some, so dropped layers retain pruned per-layer widths. Reattach the full layer list before sample(max) so every layer -- including previously-dropped ones -- is reset to the max subnet instead of being reattached in a stale state. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Sanity check on Qwen3.6-35B-A3B with MMLU 5-shot (LM Eval Harness)Original Qwen3.6-35B-A3B (34.7B, A3.5)
Pruned (22B, A3.0): Num experts 256->232, MoE FFN 512->384Pruned 37% total and 15% active params
Distilled (100 iters @ 2.5B tokens using Nemotron-Pretraining-SFT-v1)
|
|
|
||
| End-to-end tutorials that combine ModelOpt optimization techniques on [NVIDIA Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) models. | ||
| Each one walks through a complete workflow using the scripts in [examples/megatron_bridge](../README.md) (`prune_minitron.py`, `distill.py`, `quantize.py`, `export.py`). | ||
| Each one walks through a complete workflow using the scripts in [examples/megatron_bridge](../README.md) (`prune_minitron.py`, `distill.py`, `quantize.py`, `export_quantized_megatron_to_hf.py`, `export_distilled_megatron_to_hf.py`). |
There was a problem hiding this comment.
How about having one export script for any kind of megatron checkpoint, distilled or quantized or just pruned? - better UX
There was a problem hiding this comment.
Quantization export will be dropped soon. MBridge has some initial support to export quantized megatron to HF. I plan to extend it and drop custom quantized megatron import/export logic in modelopt. Leaving as it is for now
| ) | ||
| # Use object.__setattr__ to bypass DistillationProvider.__setattr__, which mirrors every attribute | ||
| # set onto the teacher -- assigning ``__class__`` normally would also switch the teacher's class. | ||
| object.__setattr__(provider, "__class__", submodule_cls) |
There was a problem hiding this comment.
it seems hacky, likely this code should be in mbridge (modifying base distillation provider there), and schedule the work, at least put TODO here to refactor it.
There was a problem hiding this comment.
Everything in this file is already in MBridge but would be available in 26.08 container. Meanwhile this file adds temporary patching to unblock VLM models. Will be removed in next mbridge release
| @@ -72,6 +72,11 @@ | |||
| from modelopt.torch.utils.plugins.megatron_mmlu import megatron_mmlu | |||
There was a problem hiding this comment.
is changing this file needed for QAD support for VLMs?
There was a problem hiding this comment.
Some bug fixes I discovered when testing Qwen3.6 pruning that was previously not tested properly. Needed to pruned + distill
| f"\nSaved quantized model to {args.export_megatron_path} in Megatron format. " | ||
| "To deploy this model (TensorRT-LLM / vLLM / SGLang), convert it to a Unified HF ckpt with export.py" | ||
| ) | ||
| if is_vlm: |
There was a problem hiding this comment.
similar comment as in other places to avoid if else paths for is_vlm
There was a problem hiding this comment.
The branch selects VLM-specific calibration (image-text dataset, calibrating only the language model) vs. text calibration -- inherent to the two modalities. Collapsing it would require unifying the LLM/VLM calibration paths, a larger refactor beyond this PR.
…ider Reframe --student_hf_model as a reference HF model with a homogeneous architecture used as the export template for a heterogeneous (Puzzletron/NAS) student (arg help in distill.py / export_distilled_megatron_to_hf.py, the export_llm_to_hf docstring, and the README). Rename the distill.py local run_provider -> distill_provider for clarity. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…t TODO Per review follow-ups: document the vision_config / language_model VLM-detection convention in distill.py, and add a TODO in export_llm_to_hf to unify the LLM export onto the in-memory path used for VLMs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
4 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