Add HybridModel MBridge support for nemo:26.08

CHANGELOG.rst

@@ -22,6 +22,7 @@ Experimental
 - Renamed ``examples/llm_ptq`` to ``examples/hf_ptq`` to reflect that it covers Hugging Face LLM **and** VLM PTQ. A relative symlink ``examples/llm_ptq`` -> ``hf_ptq`` keeps existing paths and commands working; it will be removed in a future release. Please update references to the new ``examples/hf_ptq`` path.
 - Consolidated ``examples/vlm_ptq`` into ``examples/hf_ptq``. Vision-language model PTQ now shares the ``hf_ptq.py`` entry point and ``scripts/huggingface_example.sh``; pass ``--vlm`` to run the TensorRT-LLM multimodal quickstart smoke test. The ``examples/vlm_ptq/scripts/huggingface_example.sh`` entry point is deprecated: it now prints a warning and forwards to the ``hf_ptq`` script with ``--vlm``, and will be removed in a future release. See `examples/hf_ptq/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/hf_ptq#vlm-quantization>`__.
 - Dropped VILA / NVILA vision-language model support in ``examples/hf_ptq``. VILA's modeling code requires ``transformers<=4.50.0``, which conflicts with ModelOpt's minimum supported ``transformers`` version. The VILA-specific bootstrap (repo clone, ``requirements-vila.txt``) and loading paths in ``example_utils.py`` have been removed.
+- Bump minimum nemo container requirement to ``nemo:26.04`` (recommended ``nemo:26.06``) for Megatron-Bridge / Megatron-LM optimization features.
 
 **New Features**
 

examples/megatron_bridge/distill.py

@@ -40,6 +40,7 @@
     RNGConfig,
     TokenizerConfig,
     TrainingConfig,
+    ValidationConfig,
 )
 from megatron.bridge.training.distill import distill
 from megatron.bridge.training.post_training.checkpointing import has_modelopt_state
@@ -325,15 +326,12 @@ def _restore_student_hook(model_chunks):
         model=distill_provider,
         train=TrainingConfig(
             train_iters=args.train_iters,
-            eval_interval=args.eval_interval,
-            eval_iters=args.eval_iters,
             global_batch_size=args.gbs,
             micro_batch_size=args.mbs,
             manual_gc=True,
             manual_gc_interval=100,
         ),
-        # TODO: Replace validation args in train with validation config once we drop nemo:26.02 container support
-        # validation=ValidationConfig(eval_interval=args.eval_interval, eval_iters=args.eval_iters),
+        validation=ValidationConfig(eval_iters=args.eval_iters, eval_interval=args.eval_interval),
         optimizer=optimizer_config,
         scheduler=scheduler_config,
         ddp=DistributedDataParallelConfig(

examples/megatron_bridge/prune_minitron.py

@@ -46,6 +46,15 @@
 import torch
 from megatron.bridge import AutoBridge
 from megatron.bridge.models.mamba.mamba_provider import MambaModelProvider
+
+try:  # nemo:26.08+
+    from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider
+
+    # MambaModelProvider subclasses HybridModelProvider on nemo:26.08+, so the tuple covers both.
+    _HYBRID_PROVIDER_TYPES: tuple[type, ...] = (MambaModelProvider, HybridModelProvider)
+except ImportError:  # nemo:26.06 and earlier
+    _HYBRID_PROVIDER_TYPES = (MambaModelProvider,)
+
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
@@ -541,7 +550,7 @@ def score_func(m):
         mto.ModeloptStateManager.remove_state(language_model)
     if is_vlm:
         _log_vlm_param_breakdown(unwrapped_model, language_model, "after pruning")
-    if isinstance(provider, MambaModelProvider):
+    if isinstance(provider, _HYBRID_PROVIDER_TYPES):
         hybrid_key = (
             "hybrid_override_pattern"
             if hasattr(unwrapped_model, "hybrid_override_pattern")
@@ -615,6 +624,8 @@ def score_func(m):
             if sorted_layers is not None
             else set(range(1, mcore_cfg.num_layers + 1))
         )
+        # layer_types is the HF per-layer attention-cadence field (mcore's linear_attention_freq /
+        # moe_layer_freq have no HF equivalent under those names, so only layer_types needs slicing).
         if hasattr(text_cfg, "layer_types"):
             text_cfg.layer_types = [
                 lt for i, lt in enumerate(text_cfg.layer_types) if i + 1 in kept_layer_nums
@@ -635,7 +646,9 @@ 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, MambaModelProvider) and hasattr(hf_cfg, "hybrid_override_pattern"):
+        if isinstance(provider, _HYBRID_PROVIDER_TYPES) and hasattr(
+            hf_cfg, "hybrid_override_pattern"
+        ):
             hf_cfg.hybrid_override_pattern = getattr(unwrapped_model, hybrid_key)
         text_cfg.num_hidden_layers = mcore_cfg.num_layers
         # Mark MTP as disabled on the HF text config written after pruning

examples/megatron_bridge/quantize.py

@@ -124,7 +124,7 @@ def get_args() -> argparse.Namespace:
     parser.add_argument(
         "--quant_cfg",
         type=str,
-        default="fp8",
+        default=None,
         help=(
             f"Quantization config. Preset names / short aliases: {', '.join(QUANT_CFG_CHOICES)}. "
             "You can also pass any full config name exposed by modelopt (e.g. FP8_DEFAULT_CFG). "

modelopt/torch/nas/plugins/megatron.py

@@ -88,6 +88,9 @@
 # returns the first match in insertion order: MambaModel is registered first, so
 # MambaModel instances dispatch to MambaModel whether or not MambaModel overrides forward.
 try:
+    from megatron.core.models.hybrid.hybrid_layer_specs import (
+        hybrid_stack_spec as _te_hybrid_stack_spec,
+    )
     from megatron.core.models.hybrid.hybrid_model import HybridModel
 
     SUPPORTED_MODELS[HybridModel] = "megatron.core.models.hybrid.HybridModel"
@@ -99,11 +102,11 @@
 # Attention module types that _DynamicTransformerLayer converts.
 _ATTENTION_TYPES: tuple[type, ...] = (SelfAttention, MLASelfAttention, GatedDeltaNet)
 
-__all__ = ["get_te_mamba_stack_spec"]
+__all__ = ["get_te_hybrid_stack_spec", "get_te_mamba_stack_spec"]
 
 
 def get_te_mamba_stack_spec(moe_grouped_gemm: bool = False) -> ModuleSpec:
-    """Return the TE Mamba stack spec."""
+    """[Deprecated] Return the TE Mamba stack spec."""
     assert HAS_MAMBA
     if moe_grouped_gemm:
         return _te_mamba_stack_spec
@@ -118,6 +121,21 @@ def get_te_mamba_stack_spec(moe_grouped_gemm: bool = False) -> ModuleSpec:
     return te_mamba_stack_spec
 
 
+def get_te_hybrid_stack_spec(moe_grouped_gemm: bool = False) -> ModuleSpec:
+    """Return the TE Hybrid stack spec."""
+    assert HAS_HYBRID
+    if moe_grouped_gemm:
+        return _te_hybrid_stack_spec
+
+    # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE.
+    # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency).
+    te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec)
+    te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec(
+        use_te=True, num_experts=8, moe_grouped_gemm=False
+    )
+    return te_hybrid_stack_spec
+
+
 # Local Parallel Linear DynamicModules ##########################################################################
 class _DynamicParallelLinear(DynamicModule):
     """A parallel linear layer with dynamic hyperparams."""
@@ -1055,8 +1073,12 @@ def __getattribute__(self, name):
             return mixer.d_inner
         if name in ("nheads_local_tp", "nheads_local_tpcp"):
             return mixer.nheads
-        if name == "conv1d_cp1":
+        if name == "conv1d_cp1":  # nemo:26.06 and earlier: conv is a module
             return mixer.conv1d
+        if name == "conv1d_weight_cp1":  # nemo:26.08+: raw conv parameters (dynamically sliced)
+            return mixer.conv1d_weight
+        if name == "conv1d_bias_cp1":  # nemo:26.08+
+            return mixer.conv1d_bias
         if name == "dt_bias_cp1":
             return mixer.dt_bias
         if name == "A_log_cp1":
@@ -1122,11 +1144,19 @@ def _setup(self, *, hidden_size: TracedHp):
         DMRegistry.convert(self.in_proj, input_size=hidden_size, output_size=in_proj_output_size)
 
         conv_dim = build_concat_hp([d_inner, bc])  # z, B, C
-        DMRegistry.convert(self.conv1d)
-        self.conv1d.in_channels = conv_dim
-        self.conv1d.out_channels = conv_dim
-        ks = self.conv1d.get_hparam("kernel_size")
-        ks.choices = [ks.original]
+        if hasattr(self, "conv1d"):  # nemo:26.06 and earlier: a depthwise `nn.Conv1d` module.
+            DMRegistry.convert(self.conv1d)
+            self.conv1d.in_channels = conv_dim
+            self.conv1d.out_channels = conv_dim
+            ks = self.conv1d.get_hparam("kernel_size")
+            ks.choices = [ks.original]
+        else:  # nemo:26.08+: the conv is stored as raw parameters
+
+            def _slice_conv(mod, val, _hp=conv_dim):
+                return get_sliced_tensor_by_slices(val, [_hp.active_slice])
+
+            self._register_dynamic_attribute("conv1d_weight", _slice_conv)  # [conv_dim, 1, d_conv]
+            self._register_dynamic_attribute("conv1d_bias", _slice_conv)  # [conv_dim]
 
         if self.rmsnorm:
             DMRegistry.convert(self.norm)
@@ -1151,7 +1181,8 @@ def export(self) -> torch.nn.Module:
         """Export the dynamic module to a torch.nn.Module."""
         self.in_proj.export()
         self.out_proj.export()
-        self.conv1d.export()
+        if hasattr(self, "conv1d"):  # nemo:26.06 and earlier
+            self.conv1d.export()
         if self.rmsnorm:
             self.norm.export()
         return super().export()

modelopt/torch/nas/plugins/megatron_model_stats.py

@@ -37,11 +37,17 @@
 
 import io
 import sys
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
-from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.models.mamba.mamba_model import MambaModel
+
+try:  # nemo:26.08+
+    from megatron.core.models.hybrid.hybrid_model import HybridModel
+
+    _HYBRID_MODEL_TYPES: tuple[type, ...] = (MambaModel, HybridModel)
+except ImportError:  # nemo:26.06 and earlier
+    _HYBRID_MODEL_TYPES = (MambaModel,)
 from megatron.core.parallel_state import (
     get_expert_tensor_and_model_parallel_group,
     get_expert_tensor_parallel_rank,
@@ -56,6 +62,10 @@
 from modelopt.torch.opt.dynamic import DynamicModule
 from modelopt.torch.utils import num2hrb, print_rank_0
 
+if TYPE_CHECKING:
+    from megatron.core.models.gpt.gpt_model import GPTModel
+    from megatron.core.models.hybrid.hybrid_model import HybridModel  # noqa: TC004
+
 __all__ = [
     "mcore_memory_footprint_mb",
     "mcore_param_count",
@@ -558,8 +568,8 @@ def _attn_params(i: int) -> int:
     return total, active
 
 
-def mcore_param_count_live(model: GPTModel | MambaModel) -> int:
-    """Count parameters in a live MCore GPTModel or MambaModel (reduced across TP, EP, ETP, and PP ranks)."""
+def mcore_param_count_live(model: "GPTModel | MambaModel | HybridModel") -> int:
+    """Count parameters in a live MCore LLM model (reduced across TP, EP, ETP, and PP ranks)."""
     if isinstance(model, DynamicModule):
         raise RuntimeError(
             "mcore_param_count_live does not support DynamicModule. "
@@ -718,7 +728,7 @@ def _get(attr: str, default: Any = None) -> Any:
 
 
 def print_mcore_model_stats(
-    model: "GPTModel | MambaModel",
+    model: "GPTModel | MambaModel | HybridModel",
     label: str = "Model",
     seq_length: int = 4096,
     batch_size: int = 1,
@@ -727,15 +737,15 @@ def print_mcore_model_stats(
     """Print total params, active params, and memory footprint for an MCore model.
 
     Args:
-        model: GPTModel or MambaModel to print stats for.
+        model: MCore LLM model to print stats for.
         label: Label prefix for the output line (e.g. ``"Original"``, ``"Pruned"``).
         seq_length: Sequence length for KV-cache / Mamba-state memory estimate.
         batch_size: Batch size for KV-cache / Mamba-state memory estimate.
         dtype_bytes: Bytes per parameter for memory estimation (default: 2 for BF16).
     """
     hybrid_layer_pattern: str | None = None
     config_overrides: dict = {}
-    if isinstance(model, MambaModel):
+    if isinstance(model, _HYBRID_MODEL_TYPES):
         hybrid_key = (
             "hybrid_override_pattern"
             if hasattr(model, "hybrid_override_pattern")
@@ -744,8 +754,8 @@ def print_mcore_model_stats(
         hybrid_layer_pattern = getattr(model, hybrid_key)
         # mamba_num_heads may not be stored in config when derived from model architecture;
         # fall back to reading it from the actual layer.
-        if getattr(model.config, "mamba_num_heads", None) is None:
-            for layer in model.decoder.layers:
+        if getattr(model.config, "mamba_num_heads", None) is None:  # type: ignore[attr-defined]
+            for layer in model.decoder.layers:  # type: ignore[attr-defined]
                 if hasattr(layer, "mixer") and hasattr(layer.mixer, "nheads"):
                     config_overrides["mamba_num_heads"] = layer.mixer.nheads
                     break

modelopt/torch/prune/plugins/mcore_minitron.py

@@ -168,7 +168,7 @@ def drop_mcore_language_model_layers(model: nn.Module, *, layers_to_drop: list[i
     assert isinstance(model, supported_model_types), (
         f"Model should have one of {supported_model_types} submodule, got {model}"
     )
-    print_rank_0(f"Dropping decoder layers {layers_to_drop} from model.")
+    print_rank_0(f"Dropping decoder layers {layers_to_drop} (1-indexed) from model.")
 
     # get the number of layers remaining in each pp rank
     layers_remaining_per_pp = torch.zeros(

modelopt/torch/utils/dataset_utils.py

@@ -763,7 +763,7 @@ def get_dataset_dataloader(
         batch_size: Batch size of the returned dataloader.
         num_samples: Number of samples from the dataset (interpreted as number of *output
             rows* in both ``pack=False`` and ``pack=True`` modes — in packed mode the
-            loader oversamples raw text 4x to ensure enough docs to fill all rows).
+            loader oversamples raw text 16x to ensure enough docs to fill all rows).
         max_sample_length: Maximum length of a sample (or per-row length under ``pack=True``).
         device: Target device for the returned dataloader.
         include_labels: Whether to include labels in the dataloader (ignored when
@@ -837,9 +837,9 @@ def get_dataset_dataloader(
 
     # Sample count semantics:
     # - pack=False: gather exactly `num_sample` raw docs per source, one per output row.
-    # - pack=True:  oversample 8x per source to ensure enough raw docs to fill all rows,
+    # - pack=True:  oversample 16x per source to ensure enough raw docs to fill all rows,
     #               since each row greedily packs multiple docs.
-    sample_multiplier = 8 if pack else 1
+    sample_multiplier = 16 if pack else 1
     all_samples = []
     for ds_name, num_sample in zip(dataset_name, num_samples):
         samples = get_dataset_samples(
@@ -862,10 +862,11 @@ def get_dataset_dataloader(
         )
         if input_ids.shape[0] < total_rows:
             warn_rank_0(
-                f"pack=True produced {input_ids.shape[0]} rows out of {total_rows} "
-                f"requested — raw text exhausted before filling all rows (8x oversample "
-                f"of num_samples was insufficient). Increase `num_samples` or shorten "
-                f"`max_sample_length`."
+                f"pack=True produced {input_ids.shape[0]} rows out of {total_rows} requested — "
+                f"raw text exhausted before filling all rows ({sample_multiplier}x oversample of "
+                "num_samples was insufficient). Shorten `max_sample_length` or supply longer, "
+                "more token-dense samples; increasing `num_samples` only helps if the source can "
+                "provide additional useful content."
             )
         if device:
             input_ids = input_ids.to(device)

modelopt/torch/utils/import_utils.py

@@ -15,6 +15,7 @@
 
 """Handles suppressing import errors for third-party modules that may or may not be available."""
 
+import sys
 from contextlib import contextmanager
 
 from .logging import warn_rank_0
@@ -32,7 +33,12 @@ def import_plugin(plugin_name, msg_if_missing=None, verbose=True, success_msg=No
             warn_rank_0(msg_if_missing)
     except Exception as e:
         if verbose:
+            # Capture the ``with import_plugin(...)`` call site so warnings point at the plugin
+            # that actually failed rather than at this helper. When ``__enter__`` runs the
+            # generator body, the stack is [0]=here, [1]=contextlib.__enter__, [2]=the caller.
+            caller = sys._getframe(2)
             warn_rank_0(
-                f"Failed to import modelopt {plugin_name} plugin due to: {e!r}. "
+                f"Failed to import modelopt {plugin_name} plugin "
+                f"(from {caller.f_code.co_filename}:{caller.f_lineno}) due to: {e!r}. "
                 "You may ignore this warning if you do not need this plugin."
             )

modelopt/torch/utils/plugins/mbridge.py

@@ -32,9 +32,20 @@
 from megatron.core.utils import unwrap_model
 from transformers import AutoTokenizer
 
-from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec
+from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec, get_te_mamba_stack_spec
 from modelopt.torch.utils import print_rank_0
 
+try:  # nemo:26.08+
+    from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider
+    from megatron.core.models.hybrid.hybrid_model import HybridModel
+
+    HAS_HYBRID = True
+except ImportError:
+    HAS_HYBRID = False
+    HybridModelProvider = None
+    HybridModel = None
+
+
 __all__ = ["load_mbridge_model_from_hf", "load_modelopt_megatron_checkpoint"]
 
 
@@ -48,9 +59,9 @@ def load_mbridge_model_from_hf(
     load_weights: bool = True,
 ) -> tuple[
     AutoBridge,
-    GPTModelProvider | MambaModelProvider,
+    GPTModelProvider | MambaModelProvider | HybridModelProvider,
     list[MegatronModule],
-    GPTModel | MambaModel,
+    MegatronModule,
     AutoTokenizer,
 ]:
     """Load a Megatron-Bridge model from HF.
@@ -88,8 +99,12 @@ def load_mbridge_model_from_hf(
     # provider (the bridge's native, possibly custom/hybrid spec reads it at build time) rather than
     # replacing the whole layer spec -- overwriting it would drop custom layers (e.g. Qwen3.5's
     # GatedDeltaNet + gated-attention or Gemma3's custom spec).
-    if isinstance(provider, MambaModelProvider):
+    if HAS_HYBRID and isinstance(provider, (HybridModelProvider)):
+        provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm)
+        provider.moe_grouped_gemm = moe_grouped_gemm
+    elif isinstance(provider, (MambaModelProvider)):  # Deprecated in favor of HybridModelProvider
         provider.mamba_stack_spec = get_te_mamba_stack_spec(moe_grouped_gemm=moe_grouped_gemm)
+        provider.moe_grouped_gemm = moe_grouped_gemm
     elif (provider.num_moe_experts or 0) > 0:
         provider.moe_grouped_gemm = moe_grouped_gemm
     provider.finalize()
@@ -100,10 +115,11 @@ def load_mbridge_model_from_hf(
     assert len(model) == 1
     unwrapped_model = unwrap_model(model[0])
     # VLMs (e.g. Qwen3-VL) wrap the language model as ``.language_model``; the pruning target is the
-    # inner GPTModel/MambaModel, but we still return the full wrapper so callers can save the VLM.
+    # inner GPTModel/MambaModel/HybridModel, but we still return the full wrapper so callers can save the VLM.
     language_model = getattr(unwrapped_model, "language_model", unwrapped_model)
-    assert isinstance(language_model, (GPTModel, MambaModel)), (
-        f"Expected a GPTModel/MambaModel (optionally wrapped as .language_model), "
+    model_types = (GPTModel, MambaModel, HybridModel) if HAS_HYBRID else (GPTModel, MambaModel)
+    assert isinstance(language_model, model_types), (
+        f"Expected a GPTModel/MambaModel/HybridModel (optionally wrapped as `.language_model`), "
         f"got {type(unwrapped_model)}"
     )