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. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (15)
✅ Files skipped from review due to trivial changes (6)
🚧 Files skipped from review as they are similar to previous changes (9)
📝 WalkthroughWalkthroughAdds VLM-aware Megatron-Bridge distillation, a new distilled checkpoint export script, updated checkpoint restore behavior, and broad renames from ChangesMegatron-Bridge distillation and export flow
Megatron-Bridge export name updates
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant distill.py
participant convert_to_distillation_provider
participant _convert_hook
participant mtd.convert
participant export_llm_to_hf
participant save_vlm_to_hf
distill.py->>convert_to_distillation_provider: build distillation provider
convert_to_distillation_provider->>convert_to_distillation_provider: append _convert_hook to _pre_wrap_hooks
distill.py->>convert_to_distillation_provider: load weights
convert_to_distillation_provider->>_convert_hook: invoke _convert_hook
_convert_hook->>mtd.convert: convert KD model
distill.py->>save_vlm_to_hf: export VLM after distillation
distill.py->>export_llm_to_hf: export LLM checkpoint after distillation
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✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1938 +/- ##
==========================================
- Coverage 77.79% 77.14% -0.66%
==========================================
Files 519 519
Lines 57926 58059 +133
==========================================
- Hits 45064 44789 -275
- Misses 12862 13270 +408
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.
|
- 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
623b6c9 to
3b7106a
Compare
2 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