From 58ad6edc5f61fee85a5e9632992952259049db24 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:15:47 +0530 Subject: [PATCH] Fix pruned-HF export fallback + add Nemotron-3.5-Lightning launcher examples (#2196) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix + new example Two related changes for the Megatron-Bridge Minitron prune/quantize launcher flows: 1. **Fix pruned-HF export crash on containers that reject config-only save.** `#2159` added a config-only HF export path gated only on `hasattr(AutoBridge, "from_hf_config")`. Some Megatron-Bridge versions (e.g. `nemo:26.04`) expose `from_hf_config` but reject a config-only `save_hf_pretrained` (`ValueError: save_hf_pretrained requires a pretrained HuggingFace model`), so `prune_minitron.py` crashed instead of using the intended dummy-model fallback. Now it attempts the config-only save and falls back to the dummy-model path on `ValueError`. 2. **Add Nemotron-3.5-Lightning-30B-A3B launcher examples** (`mbridge_prune.yaml`, `mbridge_quantize.yaml`) on `nemo:26.08`. Prune targets 3B active with an MMLU gate; quantize runs W4A16 NVFP4 4/6 PTQ via the `w4a16_nvfp4_4o6` recipe with `tp_size=1` (static-block NVFP4 MSE is unsupported with TP>1). ### Usage ```shell uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_prune.yaml --yes uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_quantize.yaml --yes ``` ### Testing Verified end-to-end on OCI-HSG: - **Nano prune (`nemo:26.04`)** — exercises the fallback path: config-only save raised the `ValueError`, the fallback caught it and exported via the dummy-model path. `mmlu_10pct_bs32 = 0.5196` (gate 0.50) PASS; vLLM gen PASS. - **Lightning prune (`nemo:26.08`)** — config-only export path: `score = 0.6000` (gate 0.58) PASS, 3.00B active params; vLLM gen PASS. - **Lightning quantize (`nemo:26.08`)** — recipe PTQ + unified-HF export; MMLU `0.7741` (gate 0.75) PASS. ### 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ❌ ### Additional Information The fallback fix addresses the `mbridge_prune` launcher CI failure introduced by #2159. ## Summary by CodeRabbit - **New Features** - Added a pruning workflow for Nemotron-3.5-Lightning-30B-A3B with calibration, quality scoring, checkpoint export, and multi-GPU generation. - Added a four-GPU NVFP4 W4A16 quantization workflow with Hugging Face conversion and MMLU evaluation. - **Bug Fixes** - Improved hybrid model export by falling back to dummy-model export for supported configuration-only export failures. - Added clearer logging and handling for supported export failures while preserving unrelated errors for investigation. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 --- examples/megatron_bridge/prune_minitron.py | 29 ++++++--- .../mbridge_prune.yaml | 56 +++++++++++++++++ .../mbridge_quantize.yaml | 62 +++++++++++++++++++ 3 files changed, 139 insertions(+), 8 deletions(-) create mode 100644 tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_prune.yaml create mode 100644 tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_quantize.yaml diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 7a9b4b8bd..fd8f39777 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -712,21 +712,34 @@ def main(args: argparse.Namespace): # Config-only bridge (hf_keys=None) keeps the embedding task when transformers' saved key # differs from the bridge mapping (NemotronH's backbone.embedding vs ...embeddings). - use_config_only_export = ( + exported_config_only = False + if ( hasattr(AutoBridge, "from_hf_config") and isinstance(provider, _HYBRID_PROVIDER_TYPES) and not is_vlm - ) - if use_config_only_export: + ): pruned_bridge = AutoBridge.from_hf_config(hf_cfg) # save_hf_pretrained reads trust_remote_code off the bridge to fetch source artifacts; # from_hf_config can't infer it since AutoConfig consumes the kwarg. pruned_bridge.trust_remote_code = args.trust_remote_code - pruned_bridge.save_hf_pretrained( - model, args.output_hf_path, source_path=args.hf_model_name_or_path - ) - else: - if isinstance(provider, _HYBRID_PROVIDER_TYPES) and not is_vlm: + try: + pruned_bridge.save_hf_pretrained( + model, args.output_hf_path, source_path=args.hf_model_name_or_path + ) + exported_config_only = True + except ValueError as e: + # nemo:26.06+ exposes from_hf_config but rejects config-only save_hf_pretrained; + # fall back to the dummy-model path below. + if "requires a pretrained HuggingFace model" not in str(e): + raise + warn_rank_0(f"Config-only HF export unsupported ({e}); using dummy-model export.") + + if not exported_config_only: + if ( + not hasattr(AutoBridge, "from_hf_config") + and isinstance(provider, _HYBRID_PROVIDER_TYPES) + and not is_vlm + ): warn_rank_0( "Megatron-Bridge lacks config-only HF export; falling back to the dummy-model " "path, which cannot round-trip a pruned native NemotronH config. Use " diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_prune.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_prune.yaml new file mode 100644 index 000000000..efda5d952 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_prune.yaml @@ -0,0 +1,56 @@ +# Nemotron-3.5-Lightning-30B-A3B (MoE) pruning to 3B active via Megatron-Bridge (4 GPUs), then vLLM gen. +# +# Slurm: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_prune.yaml --yes +# Local: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_prune.yaml hf_local=/mnt/hf-local --yes + +# NOTE: sized for fast CI; bump for production, e.g. --calib_num_samples 1024 --seq_length 8192 --top_k 10. +# May need to reduce batch size if running out of memory at large seq_length. +job_name: Nemotron-3.5-Lightning-30B-A3B_mbridge_prune +pipeline: + note: "Prune Nemotron-3.5-Lightning-30B-A3B -> 3B active (Megatron-Bridge) with MMLU gate, then vLLM gen" + + global_vars: + # Per-run scratch (fresh cicd_ dir) so each run prunes fresh + output_dir: /scratchspace/Nemotron-3.5-Lightning-30B-A3B-Pruned-A3.0B + + # 1) Prune and export as a HF checkpoint. + # --score_lower_bound fails the job if the pruned model's MMLU drops below the floor. + task_0: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + inline: >- + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/prune_minitron.py + --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --trust_remote_code + --pp_size 4 + --calib_batch_size 8 + --calib_num_samples 256 + --seq_length 512 + --prune_target_active_params 3e9 + --prune_target_params 24e9 + --prune_score_func mmlu_10pct_bs32 + --max_width_pruning 0.30 + --max_depth_pruning 0.15 + --hparams_to_skip num_attention_heads + --top_k 5 + --score_lower_bound 0.58 + --output_hf_path <> + slurm_config: &sc + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.08 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + docker_user: root + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 + + # 2) vLLM sanity generation on the pruned checkpoint. + task_1: + inline: >- + python modules/Model-Optimizer/examples/megatron_bridge/generate_vllm.py + --model <> + --trust_remote_code + --tensor_parallel_size 4 + slurm_config: + <<: *sc + ntasks_per_node: 1 diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_quantize.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_quantize.yaml new file mode 100644 index 000000000..7f9556bb3 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_quantize.yaml @@ -0,0 +1,62 @@ +# Nemotron-3.5-Lightning-30B-A3B NVFP4 (W4A16 4/6) quantization + unified-HF export via Megatron-Bridge (4 GPUs). +# +# Slurm: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_quantize.yaml --yes +# Local: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_quantize.yaml hf_local=/mnt/hf-local --yes + +# NOTE: sized for fast run; bump for production, e.g. --calib_num_samples 512 --seq_length 8192. +job_name: Nemotron-3.5-Lightning-30B-A3B_mbridge_quantize +pipeline: + note: "NVFP4 W4A16 PTQ Nemotron-3.5-Lightning-30B-A3B (Megatron-Bridge recipe), then unified-HF export" + + global_vars: + # Per-run scratch (fresh cicd_ dir) so each run quantizes fresh. + output_dir: /scratchspace + + # 1) NVFP4 Quantize via the ptq recipe and export to a deployable unified-HF checkpoint. + # tp_size=1: static-block NVFP4 (MSE) weight quant is unsupported with TP>1. + task_0: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + inline: >- + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/quantize.py + --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --trust_remote_code + --tp_size 1 + --recipe huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6 + --calib_batch_size 8 + --calib_num_samples 256 + --seq_length 512 + --skip_generate + --export_megatron_path <>/Nemotron-3.5-Lightning-30B-A3B-NVFP4-megatron + && + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py + --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + --megatron_path <>/Nemotron-3.5-Lightning-30B-A3B-NVFP4-megatron + --trust_remote_code + --pp_size 4 + --export_unified_hf_path <>/Nemotron-3.5-Lightning-30B-A3B-NVFP4-hf + slurm_config: &sc + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.08 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + docker_user: root + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 + + # 2) MMLU (10% sample) on the exported NVFP4 checkpoint via vLLM, gated on a lower bound. + task_1: + reqs_file: modules/Model-Optimizer/examples/llm_eval/requirements.txt + inline: >- + python modules/Model-Optimizer/examples/llm_eval/lm_eval_hf.py + --model vllm + --model_args pretrained=<>/Nemotron-3.5-Lightning-30B-A3B-NVFP4-hf,tensor_parallel_size=4 + --trust_remote_code + --tasks mmlu + --limit 0.1 + --batch_size auto + --output_path /scratchspace/mmlu_results + --accuracy_lower_bound 0.75 + slurm_config: + <<: *sc + ntasks_per_node: 1