support modelopt W4A16 & W4A4 NVFP4 export for grouped MoE weights

[HollowMan6]

What does this PR do ?

support modelopt W4A16 & W4A4 NVFP4 export for grouped MoE weights

Changelog

• Extend AutoBridge.export_hf_weights_modelopt to support quant_mode="w4a16_nvfp4". NVIDIA/Model-Optimizer@a451a2b
• Collect ModelOpt quant metadata from the actual quantized weight/module, including NVFP4 weight_scale_2.
• Map Megatron ModelOpt metadata onto exported HF names before quantization.
• Support grouped MoE expert exports by syncing EP metadata, stacking per-expert metadata, and emitting vLLM ModelOpt fused-MoE tensor names:
  • *.experts.w13_weight
  • *.experts.w13_weight_scale
  • *.experts.w13_weight_scale_2
  • *.experts.w2_weight
  • *.experts.w2_weight_scale
  • *.experts.w2_weight_scale_2
• Keep ignored weights unquantized and skip quantizer tensors during export.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[Copilot]

Pull request overview

This PR extends Megatron-Bridge’s ModelOpt export path to support W4A16 NVFP4 and improves NVFP4 export for MoE grouped expert weights by collecting richer quant metadata (including weight_scale_2), syncing metadata across EP/PP groups, and emitting vLLM fused-MoE tensor names during export.

Changes:

• Added W4A16 NVFP4 support to ModelOpt export (quant_mode="w4a16_nvfp4") and updated exporter selection logic.
• Reworked quant-metadata collection/mapping to pull metadata from the actual quantized module/weight (including weight_scale_2) and to support grouped MoE expert metadata stacking + EP syncing.
• Updated unit tests to cover new metadata mapping/quantization behaviors and adjusted patching for is_quantized.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
tests/unit_tests/models/test_modelopt_utils.py	Adds/updates unit tests for NVFP4/W4A16 NVFP4 export, metadata collection/mapping, grouped MoE handling, and fused name emission.
tests/unit_tests/models/test_auto_bridge.py	Updates mocks/patches to align with lazy is_quantized import path.
src/megatron/bridge/models/conversion/modelopt_utils.py	Implements quant metadata collection (incl. weight_scale_2), grouped expert metadata stacking, fused MoE export naming, and W4A16 NVFP4 support.
src/megatron/bridge/models/conversion/auto_bridge.py	Extends ModelOpt export flow to support new quant modes, EP+PP metadata sync, and updated metadata lookup behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[zhangyimi]

LGTM

[claude]

Light review - LGTM

Thorough ModelOpt NVFP4 / W4A16-NVFP4 export change for Nemotron-H MoE with pre-EP fused-expert quantization. The distributed logic (PP/EP metadata sync, shared-group fallback, byte-view EP all-gather, incomplete-family rejection) is intricate but matched by strong, targeted unit test coverage. Spot checks passed:

• All referenced symbols resolve (_get_ep_group, hf_to_megatron_lookup, _count_wildcard_groups, _resolve_names, extract_expert_number_from_param, NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer).
• Removed build_hf_to_megatron_name_map has no stale references in src/, docs/, or scripts/.
• Tied-embedding emit_lm_head refactor preserves prior semantics (state.source keys vs. PretrainedConfig default).
• New test_model_bridge.py carries the correct 2026 NVIDIA copyright header.
• _stage_export_tensor_to_cpu synchronizes the copy stream before returning, preserving the cpu=True contract.

Minor observations (non-blocking):

• _get_or_create_mapping reads getattr(self, attr) for the process-group attrs without a hasattr guard (unlike _copy_mapping_process_groups); safe today because MegatronParamMapping.init always sets those four attrs, but a guard would harden it against future subclasses that skip super().init.
• export_hf_weights_modelopt grew quite long; consider extracting the pre-EP eligibility/dedup block into a helper in a follow-up for readability.

Suggested test cases

• No perf tests impacted (no scripts/performance/configs/ or recipe files touched).
• Unit coverage already present and appropriate: tests/unit_tests/models/test_modelopt_utils.py, tests/unit_tests/models/test_model_bridge.py, tests/unit_tests/models/test_auto_bridge.py, tests/unit_tests/models/test_quant_mapping.py, tests/unit_tests/models/nemotronh/test_nemotron_h_bridge.py.

[claude]

Light review — LGTM

Read the full diff across auto_bridge.py, model_bridge.py, modelopt_utils.py, and param_mapping.py. Large but carefully structured feature (W4A16 & W4A4 NVFP4 export for grouped MoE weights) with strong unit coverage. No critical or high-severity issues found.

Correctness spot-checks that pass:

• Tied-embedding emit_lm_head refactor (model_bridge.py:636-658) preserves original branch semantics and clones the lm_head alias via clone_identity_output=True so the two exported tensors do not share storage.
• Removal of _install_pg_collection_on_mappings is compensated by per-mapping set_process_groups_from_pg_collection at every call site plus the new _process_group_children recursion into lazy/nested delegates.
• Fused-MoE name mapping disambiguates gate_up_proj vs up_proj correctly via ordered dict matching.
• Pre-EP local batching uses experts_per_rank; the EP finalize step byte-gathers into a rank-ordered buffer with local expert IDs rebased in _build_hf_modelopt_pre_ep_quant_metadata.
• W4A4 non-gated (.experts.up_proj) transport is explicitly rejected / routed back to the per-expert path.

Minor (non-blocking): the export_hf_weights_modelopt docstring lists nvfp4 and w4a16_nvfp4; confirm the W4A4 mode string is documented since the title advertises W4A4.

Suggested test cases (unit; run under uv run python -m pytest):

• tests/unit_tests/models/test_modelopt_utils.py::test_quantize_nvfp4_weight_exports_fused_moe_internal_names
• tests/unit_tests/models/test_modelopt_utils.py::test_quantize_nvfp4_weight_exports_fused_moe_w2_names_and_squeezes_scale_2
• tests/unit_tests/models/test_modelopt_utils.py::test_quantize_w4a16_weight_keeps_non_gated_fused_up_proj_transport_names_and_per_expert_scale_2
• tests/unit_tests/models/test_modelopt_utils.py::test_quantize_nvfp4_weight_rejects_non_gated_fused_w4a4_transport
• tests/unit_tests/models/test_modelopt_utils.py::test_collect_modelopt_quant_metadata_rejects_static_nvfp4
• tests/unit_tests/models/test_modelopt_utils.py::test_collect_modelopt_quant_metadata_max_reduces_scalars_per_tp_projection
• tests/unit_tests/models/test_modelopt_utils.py::test_build_hf_modelopt_pre_ep_quant_metadata_rebases_local_expert_ids
• tests/unit_tests/models/test_modelopt_utils.py::test_auto_bridge_modelopt_export_ep_gathers_qwen_input_scale_in_rank_order
• tests/unit_tests/models/test_modelopt_utils.py::test_auto_bridge_modelopt_export_rejects_incomplete_pre_ep_expert_family
• tests/unit_tests/models/test_modelopt_utils.py::test_auto_bridge_modelopt_export_keeps_regular_path_for_shared_pp_ep_group
• tests/unit_tests/models/test_model_bridge.py::test_stream_weights_megatron_to_hf_transforms_tied_aliases_independently
• tests/unit_tests/models/test_model_bridge.py::test_stream_weights_megatron_to_hf_batches_local_experts_before_modelopt_export
• tests/unit_tests/models/test_model_bridge.py::test_accumulate_grouped_export_uses_mapping_ep_group
• tests/unit_tests/models/test_param_mapping.py::TestModelOptLocalExpertMappings::test_process_group_collection_reaches_lazy_and_nested_mappings
• tests/unit_tests/models/test_param_mapping.py::TestModelOptLocalExpertMappings::test_fused_gated_expert_reorders_tp_shards_before_local_batching

No perf tests impacted.

[HollowMan6]

/ok to test 816c5bf