Fix pruned-HF export fallback + add Nemotron-3.5-Lightning launcher examples

examples/megatron_bridge/prune_minitron.py

@@ -712,21 +712,34 @@ def score_func(m):
 
         # 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 "

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_<id> 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 <<global_vars.output_dir>>
+    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 <<global_vars.output_dir>>
+      --trust_remote_code
+      --tensor_parallel_size 4
+    slurm_config:
+      <<: *sc
+      ntasks_per_node: 1

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_<id> 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 <<global_vars.output_dir>>/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 <<global_vars.output_dir>>/Nemotron-3.5-Lightning-30B-A3B-NVFP4-megatron
+      --trust_remote_code
+      --pp_size 4
+      --export_unified_hf_path <<global_vars.output_dir>>/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=<<global_vars.output_dir>>/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