Add Qwen3VL MCore Export support from PR 895#1482
Conversation
|
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 bidirectional Megatron Core ↔ Hugging Face weight mappings for Qwen3-VL, registers them as a plugin, includes tests validating import/export symmetry and prefix rules, and updates changelog and deployment docs for Qwen 3‑VL support. ChangesQwen3-VL Megatron Core Integration
🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 `@CHANGELOG.rst`:
- Line 136: Move the changelog entry "Add Megatron Core export/import mapping
for Qwen3-VL (``Qwen3VLForConditionalGeneration``) vision-language models..."
out of the 0.42 (2026-03-10) section and place it under the current
unreleased/0.45 section header in CHANGELOG.rst, preserving the existing
formatting and inline code markup; ensure you remove the duplicate from 0.42 and
verify the entry appears exactly once under the 0.45 (unreleased) section.
🪄 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: 7f30ba6d-8de4-4386-b197-7f8189e56d24
📒 Files selected for processing (5)
CHANGELOG.rstdocs/source/deployment/3_unified_hf.rstmodelopt/torch/export/plugins/mcore_common.pymodelopt/torch/export/plugins/mcore_qwen3vl.pytests/gpu_megatron/torch/export/test_mcore_qwen3vl.py
|
|
/claude review |
Claude Review SummarySmall, additive PR that clones the Qwen3 mcore mapping with Findings
Most impactful
Risk: Low-to-moderate. Code is purely additive, no existing arch behavior changes. Worst case is a broken Qwen3-VL round-trip that only manifests at runtime — which is exactly why the test placement matters. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1482 +/- ##
==========================================
+ Coverage 73.22% 77.24% +4.01%
==========================================
Files 478 479 +1
Lines 52421 52435 +14
==========================================
+ Hits 38387 40503 +2116
+ Misses 14034 11932 -2102
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:
|
Add Megatron Core export/import mapping for Qwen3-VL (Qwen3VLForConditionalGeneration). Handles the model.language_model. weight prefix and supports both dense and MoE variants. Signed-off-by: Hung-Yueh <hungyuehc@nvidia.com> mv test_mcore_qwen3vl.py to tests/gpu_megatron/torch/export/ Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
hychiang-git
force-pushed
the
hungyueh/pr-895
branch
from
2423ae7 to
a7d1170
Compare
Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
|
/claude review |
Claude Review SummaryFindings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2 The change is small and additive: it registers a new HF→MCore mapping for Notes
RiskLow. The change is isolated to an additive plugin entry and a new file; the round-trip has been validated by the author against |
Replace the hand-written dict literals in mcore_qwen3vl.py with a helper that derives the VL mapping from qwen3_causal_lm_import/export by inserting 'language_model.' after 'model.' in every prefix. lm_head. (root-level) is left unchanged. Remove TestQwen3VLvsQwen3Difference since it now tests the implementation against itself. Note visual encoder (model.visual.*) is intentionally excluded from the mapping. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
hychiang-git
force-pushed
the
hungyueh/pr-895
branch
from
e520fe3 to
80495e6
Compare
|
\ok to test dab32b4 |
…nents in hf_checkpoint_utils.py Import mapping (HF → mcore) — ✅ mcore_qwen3vl.py PTQ/quantize — ✅ existing functionality, no changes needed Export mapping (mcore → HF, LM weights) — ✅ mcore_qwen3vl.py + registered in mcore_common.py Vision weight merging — ✅ done, moved from test into export_mcore_gpt_to_hf via load_multimodal_components in hf_checkpoint_utils.py Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
…rmers_models.py Keep both our additions (get_tiny_qwen3vl, get_tiny_nemotron) and main's addition (get_tiny_deepseek_v3). Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
|
\ok to test 9425602 |
|
/claude review |
There was a problem hiding this comment.
Claude review passed — no blocking issues found.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 3 (all inline, non-blocking).
The PR is small, additive, and well-iterated. Verified:
_with_language_model_prefixcorrectly rewrites everymodel.-prefixed entry (language-model decoder) tomodel.language_model., leavinglm_head.at root — matches the Qwen3-VL HF safetensors layout per the linkedmodel.safetensors.index.json.- Registration in
all_mcore_hf_export_mapping/all_mcore_hf_import_mappingflows correctly through_populate_rule_book. The MoE arch (Qwen3VLMoeForConditionalGeneration) is intentionally and explicitly excluded with an in-file note, and the changelog blurb is correctly worded for the dense-only scope. load_multimodal_componentsparameterization is backward-compatible (LLaVA-style prefixes preserved as default); all-shard scan replaces the first-shard-only assumption, which is necessary for Qwen3-VL's vision encoder spanning multiple shards.save_pretrained's newelif self.arch == "Qwen3VLForConditionalGeneration"branch is mutually exclusive with theis_multimodal(LLaVA) branch and uses the samelayer_state_dicts[0].update(...)pattern as the existing path.- Test util's Qwen3VL imports are correctly lazy-loaded inside helper functions, avoiding collection breakage on older transformers builds.
- No public API,
modelopt_stateschema, or default behavior is changed; existing checkpoints/recipes are unaffected.
Inline suggestions (all non-blocking):
- Test forward pass exercises only
input_ids, so vision-encoder weight integrity isn't actually validated end-to-end (from_pretrainedonly warns on missing/unexpected keys). _with_language_model_prefixreconstructs mappings viatype(m)(target_name_or_prefix=..., func_kwargs=...), coupling to current subclass__init__signatures;copy.deepcopy(m)+ in-place rewrite is more robust.load_multimodal_componentsnow opens every shard's metadata even on the LLaVA default-prefix path; for very large checkpoints, consulting the index'sweight_mapfirst would short-circuit to only the relevant shards.
LGTM.
| vl_model = Qwen3VLForConditionalGeneration.from_pretrained( | ||
| tmp_export_dir, torch_dtype=torch.bfloat16 | ||
| ).cuda() | ||
| input_ids = torch.zeros(1, 4, dtype=torch.long).cuda() | ||
| with torch.no_grad(): | ||
| out = vl_model(input_ids=input_ids) | ||
| assert out.logits.shape[-1] == vl_model.config.text_config.vocab_size |
There was a problem hiding this comment.
[SUGGESTION] The forward pass uses input_ids only, so the vision encoder is not exercised. If the vision-weight merge step silently drops or corrupts model.visual.* tensors, this test still passes — from_pretrained only logs missing/unexpected keys at info level, and the language-model-only forward never touches the vision tower.
A cheap upgrade: capture the vision encoder state-dict from the saved tiny_qwen3vl model before export, then after export assert vl_model.model.visual.state_dict() matches it tensor-for-tensor (or at least byte-equal on a couple of representative weights). That would actually validate the cross-shard scan + merge end-to-end. Non-blocking.
| result[key] = type(m)( | ||
| target_name_or_prefix=prefix, func_kwargs=copy.deepcopy(m.func_kwargs) | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION] Reconstructing the mapping via type(m)(target_name_or_prefix=..., func_kwargs=...) couples this helper to the current CustomModuleMapping-subclass __init__ signatures (NameRemapping, QKVMerging, QKVSlicing, GatedMLPMerging, GatedMLPSlicing — all currently (target_name_or_prefix, func_kwargs)). If any future subclass adds another required init arg, or the qwen3 mapping ever uses a bare CustomModuleMapping entry (whose own __init__ takes func_name as the first positional), this rewrite quietly produces a malformed mapping.
copy.deepcopy(m) followed by an in-place result[key].target_name_or_prefix = prefix is more robust to future subclass changes and avoids re-routing through __init__. Non-blocking.
| all_shard_files = sorted(set(safetensors_index["weight_map"].values())) | ||
| first_shard_file = all_shard_files[0] # e.g., "model-00001-of-00050.safetensors" | ||
|
|
||
| # Load multimodal components from the first shard file | ||
| safetensors_filepath = Path(hf_checkpoint_path) / first_shard_file | ||
| print(f"Loading multimodal components from {first_shard_file}") | ||
|
|
||
| with safe_open(safetensors_filepath, framework="pt") as f: | ||
| shard_keys = list(f.keys()) | ||
| multimodal_keys_in_shard = [ | ||
| k for k in shard_keys if k.startswith(("multi_modal_projector", "vision_model")) | ||
| ] | ||
|
|
||
| if multimodal_keys_in_shard: | ||
| print( | ||
| f"Found {len(multimodal_keys_in_shard)} multimodal tensors in {first_shard_file}" | ||
| ) | ||
| for key in tqdm(multimodal_keys_in_shard, desc="Loading multimodal tensors"): | ||
| multimodal_state_dict[key] = f.get_tensor(key) | ||
| else: | ||
| print(f"No multimodal components found in {first_shard_file}") | ||
| for shard_file in all_shard_files: | ||
| safetensors_filepath = Path(hf_checkpoint_path) / shard_file | ||
| with safe_open(safetensors_filepath, framework="pt") as f: | ||
| for key in f.keys(): # noqa: SIM118 | ||
| if key.startswith(prefixes): | ||
| multimodal_state_dict[key] = f.get_tensor(key) |
There was a problem hiding this comment.
[SUGGESTION] This now opens every shard via safe_open and lists keys — for the LLaVA default-prefixes path this previously short-circuited to the first shard only. For checkpoints with many shards (Llama-Next-style 50+ shard layouts) and where vision components are known to be in the first shard, this adds N file-opens of metadata work. Acceptable in practice (no tensor data is loaded), but you could short-circuit by consulting safetensors_index["weight_map"] directly to determine which shards actually contain prefix-matching keys, then only safe_open those. Non-blocking.
|
/ok to test 353bef5 |
save_pretrained: hoist is_first_stage_main_rank into a single outer guard with inner is_multimodal / Qwen3-VL branches, removing the duplicated rank condition. Fix layer_state_dicts[0] -> first decoder-layer key: the dict is keyed by Megatron's 1-indexed layer_number (keys 1..num_layers), so index 0 never exists. This raised KeyError: 0 when exporting Qwen3-VL and was a latent bug in the LLaVA branch (untested). Resolve the first shard key with next(iter(layer_state_dicts)) so the vision/multimodal weights land in a shard the index builder scans (1..num_layers). test: assert every language-model decoder layer is present in the export instead of any() so a dropped-layer regression is caught. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
…nto hungyueh/pr-895
|
/ok to test 7be3e72 |
…#1792) ### What does this PR do? Type of change: New feature Adds **vision-language model (VLM) support** to the Megatron-Bridge examples for both **Minitron pruning** (`prune_minitron.py`) and **PTQ** (`quantize.py`). Only the **language model** is pruned/quantized — the vision tower and vision→language projector are left in full precision — and the full VLM is saved back. `hidden_size` is skipped for pruning when it is shared with the vision→LM projector. Supported VLMs (tested e2e): **Qwen{3,3.5}-VL** (dense; hybrid GatedDeltaNet + gated attention) and **Gemma3-VL** (sliding/full attention). ### Calibration (image-text) Calibration is conditioned on real **image-text** data so the language model's pruning importance / quantizer statistics see vision-conditioned activations. The modality is inferred from `--calib_dataset_name`: - an **image-text** dataset (default for VLMs, `nemotron_vlm_dataset_v2`) drives the **full VLM forward**; - a **text** dataset runs text-only calibration of the language model (for text-vs-image ablations). A shared `get_megatron_vlm_calibration_forward_loop` (built on `megatron_prefill`) drives the full VLM forward over image-text pairs from `vlm_dataset_utils` (`scienceqa`, `nemotron_vlm_dataset_v2`, with config-driven subset/shard caps to bound downloads). It shards across **data-parallel (DP)** ranks like the text loop (#1804); **context parallelism (CP)** applies to text-only VLM calibration (the shared text loop), not the multimodal forward — splitting the sequence would misalign the merged vision embeddings. ### Results - Cosmos-Reason2-2B Validated end-to-end on **Cosmos-Reason2-2B** (Qwen3-VL). Minitron NAS prunes the language-model tower **1.72B → ~1.59B** (vision encoder + projector frozen), top_k=1. Calibration data drives pruning importance; image-text calibration runs the full VLM forward. | Model | Calibration | MMLU | BLINK Rel-Depth | RealWorldQA | |---|---|---|---|---| | Baseline (1.72B) | — | 0.58 | 0.76 | 0.61 | | Pruned (1.59B) | text (`nemotron-post-training-dataset-v2`) | 0.51\* | ~0.69 | ~0.57 | | Pruned (1.59B) | image+text (`nemotron_vlm_dataset_v2`) | 0.49\* | **0.77** | **0.61** | \* Pruned MMLU on the 10% split (the pruning score function); baseline MMLU is the full set. The VLM-benchmark numbers for the text row were measured with a different text calibration set and are expected to be similar for `nemotron-post-training-dataset-v2` (marked `~`). > [!NOTE] > These numbers come from short single runs on small eval splits — read them for **high-level trends only**, not as exact values. Takeaways: pruning the LM tower of a VLM works end-to-end. **Image-text calibration** (this PR's feature) preserves the VLM benchmarks better than text-only — BLINK Rel-Depth ~0.77 vs ~0.69 and RealWorldQA ~0.61 vs ~0.57, both close to the unpruned baseline (0.76 / 0.61) — which is the motivation for calibrating on vision-conditioned activations. ### Results - Qwen3.5-9B | Model | MMLU | MMStar | |----------------------------|:------:|:------:| | Qwen3.5-9B | 0.7003 | 0.6117 | | Pruned-7B (text calib) | 0.5527 | 0.4411 | | Pruned-7B (image+text calib) | 0.5107 | 0.3941 | ### Key changes - `quantize.py`: quantizes the **root** model with non-LM (vision) quantizers disabled, so the ModelOpt state lives on the root (required by the Megatron save) while only the language model is quantized. - `prune_minitron.py`: image-text (or text) calibration for VLM pruning importance. - Shared VLM calibration forward loop (`megatron_prefill`-based, unwraps tuple outputs, DP-sharded) + `vlm_dataset_utils`. - Tiny VLM test fixtures (Qwen3.5-VL, Gemma3-VL) with vision tokens derived dynamically from the reference processor; VLM prune + quantize example tests. - README + CHANGELOG. ### Usage ```bash # Prune the language model of a VLM (image-text calibration by default) torchrun --nproc_per_node 2 prune_minitron.py \ --pp_size 2 \ --hf_model_name_or_path <vlm> \ --prune_target_params 3e9 \ --output_hf_path /tmp/vlm-pruned # PTQ the language model of a VLM torchrun --nproc_per_node 2 quantize.py \ --hf_model_name_or_path <vlm> \ --quant_cfg fp8 \ --export_megatron_path /tmp/vlm-fp8-megatron ``` ### Testing - `test_prune_minitron.py::test_prune_minitron_vlm` — Gemma3-VL, image-text (ScienceQA) calibration; full load → prune (depth + ffn) → save → reload. - `test_quantize_export.py::test_quantize_vlm` — Qwen3.5-VL, text calibration; quantize LM → save Megatron checkpoint. - LM regression tests (`test_prune_minitron`, `test_quantize_and_export`) unchanged and passing. ### Not in scope - **HF unified export of a quantized VLM** is not yet supported; `export.py` saves the Megatron checkpoint only for VLMs (tracked by a TODO in `export.py`). The recommended path is to route the megatron→HF quant export through Megatron-Bridge's `AutoBridge.export_hf_weights_quant(quantization_checker, quant_fn, quant_block_size)`, which reuses the bridge's per-model mcore↔HF mapping — covering Qwen3.5-VL / Gemma3-VL and the vision tower/projector (left full precision) for free — so modelopt supplies only the checker + pack/scale fn + `hf_quant_config` (KV-cache scales need a separate path). This avoids re-authoring per-model mappings in modelopt (cf. #1482's Qwen3-VL-only `mcore_qwen3vl.py`). > [!NOTE] > Qwen3.5-VL **MoE** is not tested e2e: the Megatron-Bridge weight conversion expects packed (`gate_up_proj`) experts that transformers' tiny checkpoint doesn't emit. MoE pruning itself is covered by `test_mcore_qwen35_gdn_moe_pruning`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ✅ ### Additional Information Follow-up to the GatedDeltaNet/MLA/latent-MoE pruning PR (#1747). Rebased on `main` to pick up CP/DP calibration (#1804); the VLM calibration loop now shards across DP ranks the same way. `hidden_size` pruning for VLMs (requires resizing the vision projector) is left for a future PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added VLM-aware Minitron pruning and post-training quantization that target only the language-model portion, keeping the vision tower/projector in full precision. * Calibration now auto-selects text vs image-text datasets based on model type, with modality validation. * Expanded Megatron-Core CP/DP guidance and introduced a `--cp_size` flag in quantization examples. * **Bug Fixes** * Improved VLM generation/prefill output handling and made vocabulary sizing more robust for VLM wrappers. * **Tests / Documentation** * Updated pruning/quantization docs and refreshed/added VLM-focused tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
3 participants
[Megatron Export] Add Qwen3-VL mcore ↔ HF weight mapping
What does this PR do?
New feature: Add Qwen3-VL (Vision-Language) model support to the Megatron Core export/import
plugin, enabling HuggingFace-to-mcore weight conversion for PTQ/QAT/QAD workflows.
Overview
Qwen3-VL has a different weight structure from Qwen3 text-only models:
model.language_model.prefix (notmodel.)model.visual.prefixlm_headis at root level, not nested underlanguage_modelWhat changed
modelopt/torch/export/plugins/mcore_qwen3vl.pymodel.*→model.language_model.*on top of the existing Qwen3 dense rules;lm_head.is intentionally left unchangedmodelopt/torch/export/plugins/mcore_common.pyQwen3VLForConditionalGenerationinall_mcore_hf_export_mappingandall_mcore_hf_import_mappingmodelopt/torch/export/plugins/hf_checkpoint_utils.pyload_multimodal_componentswith aprefixesparameter; sharded checkpoints now scan all shards (not just the first)modelopt/torch/export/unified_export_megatron.pysave_pretrained: added Qwen3-VL branch that copiesmodel.visual.*vision-encoder weights from the original HF checkpoint into the exported directory, producing a complete, loadable checkpointtests/_test_utils/torch/transformers_models.pyget_tiny_qwen3vl/create_tiny_qwen3vl_dirhelpers; Qwen3VL classes are lazy-imported inside the function to avoid collection failures on older transformers buildstests/gpu_megatron/torch/export/test_unified_export_megatron.pytest_unified_export_megatron/test_unified_import_megatronparametrized suites; removed standalonetest_mcore_qwen3vl.pydocs/source/deployment/3_unified_hf.rstWorkflow coverage
hf_ptqplugins/mcore_qwen3vl.py(weight name mapping),unified_export_megatron.py(export path)plugins/hf_checkpoint_utils.py(load_multimodal_componentswithprefixes),unified_export_megatron.py(calls it whenarch == "Qwen3VLForConditionalGeneration")plugins/mcore_qwen3vl.py(same mapping, reverse direction),unified_export_megatron.py(import path)Design notes
Qwen3VLMoeForConditionalGenerationstores expert weights as3-D tensors (
mlp.experts.gate_up_proj,mlp.experts.down_proj) that require adedicated fused-expert mapping. A
NotImplementedErrorcomment in the plugindocuments this explicitly.
copy.deepcopyonfunc_kwargs: each mapping entry gets its own copy toprevent shared-dict mutation when both Qwen3 and Qwen3-VL rules are loaded.
prefixesparameter onload_multimodal_components: backward-compatible defaultpreserves existing LLaVA behaviour (
"multi_modal_projector","vision_model");Qwen3-VL callers pass
("model.visual.",).Qwen3-VL vision encoder can span multiple shards, so all shards are now scanned.
Usage
From the Megatron-LM PR comment:
Create
Megatron-LM/examples/post_training/modelopt/conf/Qwen/Qwen3-VL-8B-Instruct.sh:Import Qwen3-VL from HuggingFace to MCore (local, requires GPUs):
Quantize (PTQ via Megatron-LM path):
Testing
model.language_model.prefix,lm_head.at root,QKVMerging,GatedMLPMerging,REPLICATEfor layernorms, TP sharding configsQKVSlicing,GatedMLPSlicing, noparallel_configlanguage_model.prefix,lm_headunchangedBefore your PR is "Ready for review"
tests/gpu_megatron/torch/export/test_unified_export_megatron.pydocs/source/deployment/3_unified_hf.rstCHANGELOG.rstAdditional Information
Companion Megatron-LM PR adds
Qwen3VLModel,Qwen3VLDataset, andpretrain_qwenvl.py.See: NVIDIA/Megatron-LM#3444