[kevalmorabia97]

What does this PR do?

Modernize examples/post_training/modelopt/ to use ModelOpt's plugin utilities, Adds a small megatron/core/ refactor so pruning can use the default full Transformer Engine spec without moe_grouped_gemm.

Summary

Pruning (prune.py, prune.sh, README.md)

• Single --prune-export-config '<json>' arg replaces the nine --target-* flags. --layers-to-drop removed (use num_layers in the export config).
• Calibration forward loops use modelopt.torch.utils.plugins.megatron_prefill(..., skip_return_logits=True); prompt verification uses megatron_generate.
• Pack calibration samples into uniform calib_max_sequence_length chunks instead of relying on get_dataset_dataloader's truncate-and-pad path (eliminates padding-token noise / discarded long-document context; +5.8 pts MMLU vs. unpacked on Qwen3-8B 30L/3584/11776). Revert once 0.45 ships with the upstream fix.
• Load fused TE layer_norm_weight from HF post-import for GPT-family models (_load_fused_norms_from_hf_war). Modelopt ≤ 0.44 only handles fused norms for Nemotron-H.
• Keep _gated_mlp_merging_war (modelopt ≤ 0.44 strict-load fix).
• Default calib dataset → nemotron-post-training-dataset-v2 (gated; richer multi-domain signal than the previous cnn_dailymail).
• Force args.moe_grouped_gemm = False; default --prune-intermediate-ckpt to <save>/modelopt_pruning_scores; prune.sh switches --export-te-mcore-model → --export-default-te-spec.

MMLU evaluation (mmlu.py)

• Use modelopt.torch.utils.plugins.megatron_mmlu (batched prefill + log-likelihood scoring; much faster than the previous diskcache + per-subject autoregressive loop).
• Adds --few-shots, --mmlu-batch-size. --disable-tqdm / --mmlu-dataset kept as silent BW-compat no-ops.
• WAR for the broadcast_from_last_pipeline_stage non-contiguous assertion under SP + TP (fixed upstream in 0.45).

Generation / Quantization (generate.py, quantize.py, megatron/post_training/generate.py)

• Swap simple_generate → megatron_generate on all text-gen paths; calibration uses megatron_prefill(skip_return_logits=True).
• simple_speculative_generate rewired through megatron_generate(enable_kv_cache=False) — incidentally fixes a latent position_ids bug (arange(shape[0]) was using batch size).
• simple_generate deleted — megatron_generate / megatron_prefill cover every prior call site. Module loses several transitive deps.
• Default calib dataset → nemotron-post-training-dataset-v2 (matches prune.py).

Core refactor (megatron/core/, megatron/post_training/model_builder.py)

• New get_hybrid_stack_spec(moe_grouped_gemm: bool = True) factory in hybrid_layer_specs.py; module-level hybrid_stack_spec / mamba_stack_spec preserved.
• get_hybrid_stack_modelopt_spec threads moe_grouped_gemm to the factory when use_default_te_spec=True.
• modelopt_gpt_hybrid_builder's GPT branch honours --export-default-te-spec (was hybrid-only) and forwards args.moe_grouped_gemm.

Upstream consumers

WARs above are temporary; the matching modelopt 0.45 fixes are in Model-Optimizer PR #1501. Once 0.45 ships and the modelopt pin is bumped, the prune.py / mmlu.py WARs can be removed.

Pre-checks

• Typing on public API
• Documentation updated
• tools/autoformat.sh run (black + isort + ruff + pylint + mypy)
• Unit tests (n/a — example refactor + small core helper)
• Functional tests (covered by internal nmm-sandbox CI)

cc @mcore-oncall

[copy-pr-bot]

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.

[kevalmorabia97]

/ok to test 0f67913

[kevalmorabia97]

/ok to test 239971b

[kevalmorabia97]

/ok to test 5753fc1

[kevalmorabia97]

/ok to test e60e95f

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26124136128