[Cherry-pick] PRs #1975 #2076 #2071 #2093 #2084 #2115 #2133 #2146 #2064 #2159 #2112

examples/hf_ptq/README.md

@@ -118,7 +118,6 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http
 | Whisper<sup>9</sup> | ✅ | ❌ | ❌ | ❌ | - |
 | Nemotron-3 | ✅ | ❌ | ❌ | ❌ | ✅ |
 | Llava (VLM)<sup>11</sup> | ✅ | ✅<sup>12</sup> | ✅ | ✅ | - |
-| Phi-3-vision, Phi-4-multimodal (VLM)<sup>11</sup> | ✅ | ✅<sup>12</sup> | ✅ | ✅ | ✅ |
 | Qwen2, 2.5-VL (VLM)<sup>11</sup> | ✅ | ✅<sup>12</sup> | ✅ | ✅ | ✅ |
 | Gemma 3 (VLM)<sup>11</sup> | ✅ | - | - | - | - |
 | Nemotron VL (VLM)<sup>11,13</sup> | ✅ | - | - | - | ✅ |

examples/hf_ptq/example_utils.py

@@ -168,12 +168,6 @@ def _is_multimodal_config(config):
     """Check if a config indicates a multimodal model (config-only version of is_multimodal_model)."""
     return (
         hasattr(config, "vision_config")  # Standard vision config (e.g., Qwen2.5-VL)
-        or getattr(config, "model_type", "") == "phi4mm"  # Phi-4 multimodal
-        or hasattr(config, "vision_lora")  # Vision LoRA configurations
-        or hasattr(config, "audio_processor")  # Audio processing capabilities
-        or (
-            hasattr(config, "embd_layer") and hasattr(config.embd_layer, "image_embd_layer")
-        )  # Image embedding layers
         or getattr(config, "is_encoder_decoder", False)  # Encoder-decoder VL models
         or any(  # Architecture-based detection for custom VL models (e.g., Nemotron-Parse)
             "conditionalgeneration" in arch.lower() for arch in getattr(config, "architectures", [])

examples/hf_ptq/hf_ptq.py

@@ -699,9 +699,6 @@ def load_model(args: argparse.Namespace):
         # Left padding usually provides better calibration result.
         tokenizer.padding_side = "left"
 
-    if model_type == "phi4mm":
-        warnings.warn("Please set the default input_mode to InputMode.LANGUAGE before quantizing.")
-
     return (
         full_model,
         language_model,

examples/llm_qat/export.py

@@ -23,7 +23,7 @@
 import modelopt.torch.opt as mto
 from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
 from modelopt.torch.export.unified_export_hf import _export_transformers_checkpoint
-from modelopt.torch.opt.conversion import restore_from_modelopt_state
+from modelopt.torch.opt.conversion import ModeloptStateManager, restore_from_modelopt_state
 from modelopt.torch.quantization.utils import set_quantizer_state_dict
 from modelopt.torch.utils import print_rank_0
 
@@ -48,8 +48,11 @@ def get_model(
     # Load model
     model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map=device_map)
 
-    # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this
-    if hasattr(model, "peft_config"):
+    # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this.
+    # For QLoRA the base checkpoint is quantized, so from_pretrained already restored the state.
+    # Skipping is safe only because QATTrainer writes modelopt_state_train.pth at trainer init,
+    # from that same base state.
+    if hasattr(model, "peft_config") and not ModeloptStateManager.is_converted(model):
         modelopt_state = mto.load_modelopt_state(f"{ckpt_path}/modelopt_state_train.pth")
         restore_from_modelopt_state(model, modelopt_state)
         print_rank_0("Restored modelopt state")

examples/megatron_bridge/README.md

@@ -365,9 +365,6 @@ torchrun --nproc_per_node 1 prune_minitron.py --help
 > [!NOTE]
 > Multi-token-prediction (MTP) heads (e.g. Qwen3.5) are not pruned yet — they are dropped for the prune run and the saved checkpoint has no MTP. Autoregressive inference is unaffected; for speculative decoding, run a short MTP SFT on the pruned model.
 
-> [!NOTE]
-> If pruning a Nemotron model and you want to save the pruned model back in HF format, please downgrade to `transformers<5` via `python -m pip install "transformers<5"` before pruning.
-
 ### Vision-Language Models (VLMs)
 
 For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `prune_minitron.py` automatically prunes only the **language model** and leaves the vision tower intact, then saves the full VLM back. All the pruning modes above (parameter count, active parameter count, memory footprint, and manual `export_config`) work unchanged, with two VLM-specific caveats:

examples/megatron_bridge/distill.py

@@ -480,5 +480,7 @@ def _restore_student_hook(model_chunks):
     args = get_args()
     try:
         main(args)
+    except BaseException:
+        dist.abort()  # peers may be stuck in a collective this rank will never reach
     finally:
         dist.cleanup()

examples/megatron_bridge/export_distilled_megatron_to_hf.py

@@ -288,5 +288,7 @@ def main(args: argparse.Namespace):
     args = get_args()
     try:
         main(args)
+    except BaseException:
+        dist.abort()  # peers may be stuck in a collective this rank will never reach
     finally:
         dist.cleanup()

examples/megatron_bridge/export_quantized_megatron_to_hf.py

@@ -164,5 +164,7 @@ def main(args: argparse.Namespace):
     args = get_args()
     try:
         main(args)
+    except BaseException:
+        dist.abort()  # peers may be stuck in a collective this rank will never reach
     finally:
         dist.cleanup()

examples/megatron_bridge/prune_minitron.py

@@ -105,6 +105,24 @@ def _hf_config_has_mtp(hf_cfg) -> bool:
     )
 
 
+# HF names the shared expert size with or without the ``moe_`` prefix depending on the model
+# (e.g. Qwen3.5-MoE uses ``shared_expert_intermediate_size``).
+_SHARED_EXPERT_SIZE_FIELDS = (
+    "moe_shared_expert_intermediate_size",
+    "shared_expert_intermediate_size",
+)
+
+
+def _is_deepseek_style_moe(text_cfg) -> bool:
+    """Whether the shared expert is sized as ``n_shared_experts * moe_intermediate_size``.
+
+    Such configs can only represent a shared expert size that is a multiple of the routed one.
+    """
+    return getattr(text_cfg, "n_shared_experts", None) is not None and not any(
+        hasattr(text_cfg, field) for field in _SHARED_EXPERT_SIZE_FIELDS
+    )
+
+
 def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("--hf_model_name_or_path", type=str, required=True)
@@ -400,7 +418,9 @@ def main(args: argparse.Namespace):
             "num_layers_in_last_pipeline_stage": args.num_layers_in_last_pipeline_stage,
             "pipeline_dtype": torch.bfloat16,
             "seq_length": args.seq_length,
-            "mtp_num_layers": 0,  # MTP is not supported during calibration
+            # MTP is not supported during calibration; drop it
+            "mtp_num_layers": 0,
+            "mtp_hybrid_override_pattern": None,
         },
         init_model_parallel=True,
         moe_grouped_gemm=not args.no_moe_grouped_gemm,
@@ -542,6 +562,17 @@ def score_func(m):
         pruning_config["max_width_pruning"] = args.max_width_pruning
         pruning_config["max_depth_pruning"] = args.max_depth_pruning
         pruning_config["hparams_to_skip"] = args.hparams_to_skip
+        # DeepSeek-style MoE configs size the shared expert as n_shared_experts * moe_intermediate_size,
+        # so only candidates whose shared size is a multiple of the routed one can be saved to HF.
+        src_hf_cfg = bridge.hf_pretrained.config
+        if _is_deepseek_style_moe(getattr(src_hf_cfg, "text_config", src_hf_cfg)):
+            warn_rank_0(
+                "DeepSeek-style MoE config detected: restricting the search to candidates whose "
+                "moe_shared_expert_intermediate_size is a multiple of moe_ffn_hidden_size."
+            )
+            pruning_config["candidate_filter"] = lambda cfg: (
+                cfg["moe_shared_expert_intermediate_size"] % cfg["moe_ffn_hidden_size"] == 0
+            )
         pruning_config["top_k"] = args.top_k
         # memory_mb constraint requires batch_size and seq_length
         pruning_config["batch_size"] = args.inference_batch_size
@@ -588,10 +619,11 @@ def score_func(m):
     else:
         print_rank_0(f"Saving pruned model to {args.output_hf_path} in HF checkpoint format")
 
-        # [WAR] Save the pruned HF model by hand until Megatron-Bridge natively supports it.
-        # TODO: Replace this whole block with ``AutoBridge.from_auto_config(...).save_hf_weights(...)``
-        #     once the Megatron-Bridge fix ships (nemo:26.08).
-        bridge.hf_pretrained.save_artifacts(args.output_hf_path)
+        # Build the pruned HF config field-by-field from the pruned Megatron config, then stream weights.
+        # Rank 0 only: a late write from another rank would leave config.json stale.
+        if dist.is_master():
+            bridge.hf_pretrained.save_artifacts(args.output_hf_path)
+        dist.barrier()
         hf_cfg = AutoConfig.from_pretrained(
             args.output_hf_path, trust_remote_code=args.trust_remote_code
         )
@@ -610,12 +642,7 @@ def score_func(m):
             text_cfg.mamba_head_dim = mcore_cfg.mamba_head_dim
         if hasattr(text_cfg, "moe_intermediate_size"):
             text_cfg.moe_intermediate_size = mcore_cfg.moe_ffn_hidden_size
-        # HF names this field with or without the ``moe_`` prefix depending on the model
-        # (e.g. Qwen3.5-MoE uses ``shared_expert_intermediate_size``).
-        for shared_expert_field in (
-            "moe_shared_expert_intermediate_size",
-            "shared_expert_intermediate_size",
-        ):
+        for shared_expert_field in _SHARED_EXPERT_SIZE_FIELDS:
             if hasattr(text_cfg, shared_expert_field):
                 setattr(
                     text_cfg, shared_expert_field, mcore_cfg.moe_shared_expert_intermediate_size
@@ -624,7 +651,16 @@ def score_func(m):
             text_cfg.num_experts = mcore_cfg.num_moe_experts
         if hasattr(text_cfg, "n_routed_experts"):
             text_cfg.n_routed_experts = mcore_cfg.num_moe_experts
-        if hasattr(text_cfg, "n_shared_experts"):
+        # n_shared_experts is a fixed count; only DeepSeek-style configs record the pruned shared
+        # expert size through it. candidate_filter keeps the search divisible, so only a manual
+        # --prune_export_config can violate this.
+        if _is_deepseek_style_moe(text_cfg):
+            if mcore_cfg.moe_shared_expert_intermediate_size % mcore_cfg.moe_ffn_hidden_size:
+                raise ValueError(
+                    f"{mcore_cfg.moe_shared_expert_intermediate_size=} must be a multiple of "
+                    f"{mcore_cfg.moe_ffn_hidden_size=} for this config, which stores the shared "
+                    "expert size as n_shared_experts * moe_intermediate_size. "
+                )
             text_cfg.n_shared_experts = (
                 mcore_cfg.moe_shared_expert_intermediate_size // mcore_cfg.moe_ffn_hidden_size
             )
@@ -658,8 +694,11 @@ def score_func(m):
                     "distillation cannot recover this vision-path change -- consider full VLM "
                     "training/distillation instead of LM-only to recover vision quality."
                 )
-        if isinstance(provider, _HYBRID_PROVIDER_TYPES) and hasattr(
-            text_cfg, "hybrid_override_pattern"
+        # Only older remote-code configs need this; native configs carry the cadence in layer_types.
+        if (
+            isinstance(provider, _HYBRID_PROVIDER_TYPES)
+            and not hasattr(text_cfg, "layer_types")
+            and hasattr(text_cfg, "hybrid_override_pattern")
         ):
             # MCore's pattern can carry an MTP suffix (``/...``) and PP boundaries (``|``) which we need to remove
             text_cfg.hybrid_override_pattern = "".join(
@@ -671,15 +710,36 @@ def score_func(m):
             if hasattr(text_cfg, field):
                 setattr(text_cfg, field, 0)
 
-        # Save dummy pruned HF model to get the correct bridge for saving pruned weights
-        dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
-        dummy_model_cls.from_config(
-            hf_cfg, trust_remote_code=args.trust_remote_code
-        ).save_pretrained(args.output_hf_path, trust_remote_code=args.trust_remote_code)
-        pruned_bridge = AutoBridge.from_hf_pretrained(
-            args.output_hf_path, trust_remote_code=args.trust_remote_code
+        # 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 = (
+            hasattr(AutoBridge, "from_hf_config")
+            and isinstance(provider, _HYBRID_PROVIDER_TYPES)
+            and not is_vlm
         )
-        pruned_bridge.save_hf_weights(model, args.output_hf_path)
+        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:
+                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 "
+                    "transformers<5 or a newer Megatron-Bridge if the save fails."
+                )
+            dummy_model_cls = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
+            dummy_model_cls.from_config(
+                hf_cfg, trust_remote_code=args.trust_remote_code
+            ).save_pretrained(args.output_hf_path, trust_remote_code=args.trust_remote_code)
+            pruned_bridge = AutoBridge.from_hf_pretrained(
+                args.output_hf_path, trust_remote_code=args.trust_remote_code
+            )
+            pruned_bridge.save_hf_weights(model, args.output_hf_path)
 
         copy_hf_ckpt_remote_code(args.hf_model_name_or_path, args.output_hf_path)
         print_rank_0(f"Saved pruned model to {args.output_hf_path} in HF checkpoint format")
@@ -704,5 +764,7 @@ def score_func(m):
     args = get_args()
     try:
         main(args)
+    except BaseException:
+        dist.abort()  # peers may be stuck in a collective this rank will never reach
     finally:
         dist.cleanup()

examples/megatron_bridge/quantize.py

@@ -454,5 +454,7 @@ def forward_loop(_model=None):
     args = get_args()
     try:
         main(args)
+    except BaseException:
+        dist.abort()  # peers may be stuck in a collective this rank will never reach
     finally:
         dist.cleanup()

examples/speculative_decoding/eagle_utils.py

@@ -141,6 +141,9 @@ def make_speculative_data_module(
                 train_len=train_len,
                 local_image_path=data_args.vlm_img_dir,
                 return_labels=True,
+                answer_only_loss=answer_only_loss,
+                shift_labels=shift_labels,
+                chat_template=chat_template,
             )
 
     else:

examples/vllm_serve/Dockerfile

@@ -1,4 +1,4 @@
-FROM vllm/vllm-openai:v0.20.0
+FROM vllm/vllm-openai:v0.26.0
 
 # Set environment variables
 ENV PIP_NO_CACHE_DIR=off \
@@ -25,12 +25,16 @@ RUN cd Model-Optimizer && \
 # Llama4 requires this
 RUN pip install flash-attn==2.7.4.post1 --no-build-isolation
 
-# Pre-compile CUDA extensions to avoid compilation time during runtime
+# Pre-compile CUDA extensions into a world-accessible directory so the vllm
+# user can use the cache at runtime.
+ENV TORCH_EXTENSIONS_DIR=/workspace/torch_extensions
 RUN python3 -c "import modelopt.torch.quantization.extensions as ext; ext.precompile()" || true
 
-# Allow users to run without root
+# Allow the non-root vllm user to access the workspace
 RUN chmod -R 777 /workspace
 
+USER vllm
+
 # Override the ENTRYPOINT from the base image to allow flexible usage
 ENTRYPOINT []
 

examples/vllm_serve/README.md

@@ -4,7 +4,7 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa
 
 Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research.
 
-The general fakequant example is tested with vLLM 0.9.0 and 0.19.1. The compact
+The general fakequant example is tested with vLLM 0.9.0, 0.19.1, and 0.26.0. The compact
 NVFP4 attention worker documented below requires vLLM 0.15.0 or newer.
 
 ## Prepare environment

examples/vllm_serve/vllm_ptq_utils.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import dataclasses
+import warnings
 from collections.abc import Callable
 from typing import Any
 
@@ -66,6 +67,7 @@ def calibrate_loop(model: Any) -> None:
                     NewRequestData,
                     req_id=req_id,
                     prompt_token_ids=input_ids_list,
+                    prefill_token_ids=input_ids_list,
                     mm_kwargs=[],
                     mm_hashes=[],
                     mm_positions=[],
@@ -95,10 +97,39 @@ def calibrate_loop(model: Any) -> None:
                 structured_output_request_ids={},
                 grammar_bitmask=None,
             )
-            output = self.execute_model(scheduler_output)
-            if hasattr(self, "sample_tokens"):
-                if output is None:  # TODO: make this default when vllm <= 0.11 is outdated
-                    self.sample_tokens(None)
+            try:
+                output = self.execute_model(scheduler_output)
+                if hasattr(self, "sample_tokens"):
+                    if output is None:  # TODO: make this default when vllm <= 0.11 is outdated
+                        self.sample_tokens(None)
+            finally:
+                # finish_requests runs before add_requests inside execute_model, so
+                # req IDs aren't registered yet at that point — call it directly after.
+                # Wrap in try/except so a cleanup error never masks the original exception.
+                try:
+                    if hasattr(self.model_runner, "finish_requests"):
+                        cleanup_output = _create_new_data_cls(
+                            type(scheduler_output),
+                            scheduled_new_reqs=[],
+                            scheduled_cached_reqs=scheduler_output.scheduled_cached_reqs,
+                            num_scheduled_tokens={},
+                            total_num_scheduled_tokens=0,
+                            scheduled_spec_decode_tokens={},
+                            scheduled_encoder_inputs={},
+                            num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks,
+                            finished_req_ids=set(num_scheduled_tokens.keys()),
+                            free_encoder_mm_hashes=[],
+                            kv_connector_metadata=None,
+                            structured_output_request_ids={},
+                            grammar_bitmask=None,
+                        )
+                        self.model_runner.finish_requests(cleanup_output)
+                    else:
+                        warnings.warn(
+                            "model_runner.finish_requests not found; request state may leak during calibration."
+                        )
+                except Exception:
+                    warnings.warn("Failed to clean up request state after calibration batch.")
 
     return calibrate_loop
 

modelopt/torch/export/layer_utils.py

@@ -222,12 +222,7 @@ def is_conv(module: nn.Module) -> bool:
 def is_embedding(module: nn.Module) -> bool:
     """Returns whether the module is an embedding layer."""
     module_type_name = type(module).__name__
-    return (
-        "Embedding" in module_type_name
-        and "Rotary" not in module_type_name
-        and "PhiImage" not in module_type_name
-        and "Phi3Image" not in module_type_name
-    )
+    return "Embedding" in module_type_name and "Rotary" not in module_type_name
 
 
 def build_embedding_config(module: nn.Module, normalization_constant: float = 1) -> EmbeddingConfig: